import type { InstanceAiAgentNode } from '@n8n/api-types';
import { Service } from '@n8n/di';
import type { AgentTreeSnapshot } from '@n8n/instance-ai';
import { jsonParse } from 'n8n-workflow';

import { InstanceAiRunSnapshotRepository } from '../repositories/instance-ai-run-snapshot.repository';

/**
 * Walk a saved agent tree and flip everything in-flight to a terminal state.
 * Active sub-agents become `cancelled`, loading tool calls stop loading, and
 * unresolved HITL confirmation cards get a `denied` status so the frontend
 * stops rendering Allow / Request-changes buttons. Mutates `node` in place.
 */
function cancelInFlightNodes(node: InstanceAiAgentNode): void {
	if (node.status === 'active') node.status = 'cancelled';
	for (const call of node.toolCalls) {
		if (!call.isLoading) continue;
		call.isLoading = false;
		if (call.confirmation && !call.confirmationStatus) {
			call.confirmationStatus = 'denied';
		}
	}
	for (const child of node.children) cancelInFlightNodes(child);
}

export interface SaveSnapshotOptions {
	messageGroupId?: string;
	runIds?: string[];
	traceId?: string;
	spanId?: string;
	langsmithRunId?: string;
	langsmithTraceId?: string;
}

@Service()
export class DbSnapshotStorage {
	constructor(private readonly repo: InstanceAiRunSnapshotRepository) {}

	async getLatest(
		threadId: string,
		options: { messageGroupId?: string; runId?: string } = {},
	): Promise<AgentTreeSnapshot | undefined> {
		const { messageGroupId, runId } = options;

		const row = messageGroupId
			? await this.repo.findOne({
					where: { threadId, messageGroupId },
					order: { createdAt: 'DESC' },
				})
			: runId
				? await this.repo.findOne({
						where: { threadId, runId },
						order: { createdAt: 'DESC' },
					})
				: await this.repo.findOne({
						where: { threadId },
						order: { createdAt: 'DESC' },
					});

		if (!row) return undefined;

		return {
			tree: jsonParse<InstanceAiAgentNode>(row.tree),
			runId: row.runId,
			messageGroupId: row.messageGroupId ?? undefined,
			runIds: row.runIds ?? undefined,
			traceId: row.traceId ?? undefined,
			spanId: row.spanId ?? undefined,
			langsmithRunId: row.langsmithRunId ?? undefined,
			langsmithTraceId: row.langsmithTraceId ?? undefined,
			createdAt: row.createdAt,
			updatedAt: row.updatedAt,
		};
	}

	async save(
		threadId: string,
		agentTree: InstanceAiAgentNode,
		runId: string,
		options: SaveSnapshotOptions = {},
	): Promise<void> {
		const { messageGroupId, runIds, traceId, spanId, langsmithRunId, langsmithTraceId } = options;
		await this.repo.upsert(
			{
				threadId,
				runId,
				messageGroupId: messageGroupId ?? null,
				runIds: runIds ?? null,
				tree: JSON.stringify(agentTree),
				traceId: traceId ?? null,
				spanId: spanId ?? null,
				langsmithRunId: langsmithRunId ?? null,
				langsmithTraceId: langsmithTraceId ?? null,
			},
			['threadId', 'runId'],
		);
	}

	async updateLast(
		threadId: string,
		agentTree: InstanceAiAgentNode,
		runId: string,
		options: SaveSnapshotOptions = {},
	): Promise<void> {
		const { messageGroupId, runIds, traceId, spanId, langsmithRunId, langsmithTraceId } = options;

		// Prefer lookup by messageGroupId when available
		if (messageGroupId) {
			const existing = await this.repo.findOne({
				where: { threadId, messageGroupId },
				order: { createdAt: 'DESC' },
			});
			if (existing) {
				await this.repo.update(
					{ threadId: existing.threadId, runId: existing.runId },
					{
						runId,
						tree: JSON.stringify(agentTree),
						messageGroupId,
						runIds: runIds ?? existing.runIds,
						// Preserve existing trace IDs if caller didn't provide new ones.
						traceId: traceId ?? existing.traceId,
						spanId: spanId ?? existing.spanId,
						langsmithRunId: langsmithRunId ?? existing.langsmithRunId,
						langsmithTraceId: langsmithTraceId ?? existing.langsmithTraceId,
					},
				);
				return;
			}
		}

		// Fall back to runId lookup
		const byRunId = await this.repo.findOneBy({ threadId, runId });
		if (byRunId) {
			await this.repo.update(
				{ threadId, runId },
				{
					tree: JSON.stringify(agentTree),
					messageGroupId: messageGroupId ?? byRunId.messageGroupId,
					runIds: runIds ?? byRunId.runIds,
					traceId: traceId ?? byRunId.traceId,
					spanId: spanId ?? byRunId.spanId,
					langsmithRunId: langsmithRunId ?? byRunId.langsmithRunId,
					langsmithTraceId: langsmithTraceId ?? byRunId.langsmithTraceId,
				},
			);
			return;
		}

		// No existing row — insert
		await this.save(threadId, agentTree, runId, options);
	}

	/**
	 * Mark an existing snapshot as a cancelled run, terminalising every
	 * `active` node and every in-flight tool call (including unresolved HITL
	 * confirmations) in the saved tree. Used when a run is being marked
	 * terminal after the in-memory event bus is gone — e.g. handling a
	 * confirmation orphaned by a restart — because rebuilding the tree from
	 * an empty bus would clobber the saved plan / ask card with an empty
	 * cancelled tree. Keeps the planner output, tool calls, and confirmation
	 * payload intact so the user can still see what was being planned, just
	 * without interactive buttons.
	 */
	async markRunCancelled(threadId: string, runId: string): Promise<void> {
		const key = { threadId, runId };
		const row = await this.repo.findOneBy(key);
		if (!row) return;

		const tree = jsonParse<InstanceAiAgentNode>(row.tree);
		cancelInFlightNodes(tree);
		await this.repo.update(key, { tree: JSON.stringify(tree) });
	}

	async getAll(threadId: string): Promise<AgentTreeSnapshot[]> {
		const rows = await this.repo.find({
			where: { threadId },
			order: { createdAt: 'ASC' },
		});
		return rows.map((r) => ({
			tree: jsonParse<InstanceAiAgentNode>(r.tree),
			runId: r.runId,
			messageGroupId: r.messageGroupId ?? undefined,
			runIds: r.runIds ?? undefined,
			traceId: r.traceId ?? undefined,
			spanId: r.spanId ?? undefined,
			langsmithRunId: r.langsmithRunId ?? undefined,
			langsmithTraceId: r.langsmithTraceId ?? undefined,
			createdAt: r.createdAt,
			updatedAt: r.updatedAt,
		}));
	}

	/**
	 * Resolve the LangSmith root-run anchor for a given responseId
	 * (UI sends `messageGroupId ?? runId`). Prefers the earliest snapshot row
	 * in a message group so feedback attaches to the `message_turn` root run.
	 */
	async findLangsmithAnchor(
		threadId: string,
		responseId: string,
	): Promise<{ langsmithRunId: string; langsmithTraceId: string } | undefined> {
		const byGroup = await this.repo.findOne({
			where: { threadId, messageGroupId: responseId },
			order: { createdAt: 'ASC' },
		});
		const row = byGroup ?? (await this.repo.findOneBy({ threadId, runId: responseId }));
		if (!row?.langsmithRunId || !row.langsmithTraceId) return undefined;
		return { langsmithRunId: row.langsmithRunId, langsmithTraceId: row.langsmithTraceId };
	}
}