import type {
	InstanceAiCredentialDestinationDecision,
	InstanceAiThreadStatusResponse,
} from '@n8n/api-types';
import { nanoid } from 'nanoid';

import type { InstanceAiTraceContext, ModelConfig, OrchestrationContext } from '../types';
import type {
	InstanceAiLivenessPolicy,
	InstanceAiLivenessSurface,
	InstanceAiLivenessTimeoutReason,
} from './liveness-policy';
import type { OrchestratorRunHandoffState } from './orchestrator-run-control';
import type { WorkflowBuildOutcome } from '../workflow-loop/workflow-loop-state';

export interface ActiveRunState {
	runId: string;
	threadId: string;
	abortController: AbortController;
	/** Prevents an older finalizer from clearing a newer executor for the same thread. */
	executionToken?: symbol;
	messageGroupId?: string;
	tracing?: InstanceAiTraceContext;
	modelId?: ModelConfig;
	startedAt?: number;
	lastActivityAt?: number;
	/** Owner of the run, resolved at entry so concurrency can be counted per user. */
	userId?: string;
}

export interface SuspendedRunState<TUser = unknown> extends ActiveRunState {
	agentRunId: string;
	agent: unknown;
	/** The orchestration context the agent's tools closed over. Stored so a
	 *  resume can rebind `tracing` to the new resume trace — spans emitted
	 *  through the suspended turn's shut-down runtime export nothing. */
	orchestrationContext?: OrchestrationContext;
	threadId: string;
	user: TUser;
	toolCallId: string;
	toolName?: string;
	suspendPayload?: Record<string, unknown>;
	requestId: string;
	createdAt: number;
	/** Set when the suspended run was a planned-task checkpoint follow-up.
	 *  Preserved across suspend/resume so the resumed run's finalizer can
	 *  run the deadlock fallback and reschedule. */
	checkpoint?: { isCheckpointFollowUp: true; checkpointTaskId: string };
	/** Set when the suspended run was a planned build-workflow follow-up. */
	plannedBuild?: {
		isPlannedBuildFollowUp: true;
		buildTaskId: string;
		workItemId: string;
		isSupportingWorkflowTask?: boolean;
		savedOutcome?: WorkflowBuildOutcome;
	};
	/** Shared signal used to stop resumed orchestration after durable work is handed off. */
	runHandoff?: OrchestratorRunHandoffState;
}

/**
 * Flat confirmation payload consumed by native tool `resumeSchema`s and sub-agent HITL.
 * The service layer constructs this from the typed `InstanceAiConfirmRequest` discriminated
 * union sent by the frontend — only one subset of fields is populated per call, matching
 * the confirmation kind that was originally requested.
 */
export interface ConfirmationData {
	approved: boolean;
	credentials?: Record<string, string>;
	nodeCredentials?: Record<string, Record<string, string>>;
	userInput?: string;
	domainAccessAction?: string;
	action?: 'apply' | 'test-trigger';
	nodeParameters?: Record<string, Record<string, unknown>>;
	/** Workflow-setup cards the user actively skipped, by node name. */
	skippedNodes?: string[];
	testTriggerNode?: string;
	answers?: Array<{
		questionId: string;
		selectedOptions: string[];
		customText?: string;
		skipped?: boolean;
	}>;
	/** User's resource-access decision (e.g. 'allowForSession'). */
	resourceDecision?: string;
	/** Plan-review hard denial — distinct from a feedback-driven rejection. */
	denied?: boolean;
	/** `'session'` means the user chose "always allow": the resuming tool should
	 *  persist a thread-level grant so the same action isn't re-asked. */
	scope?: 'once' | 'session';
	autoSetup?: { credentialType: string; attemptId?: string };
	credentialDestination?: InstanceAiCredentialDestinationDecision;
	connectedSlugs?: string[];
}

export interface PendingConfirmation {
	resolve: (data: ConfirmationData) => void;
	threadId: string;
	userId: string;
	createdAt: number;
	startedAt?: number;
	lastActivityAt?: number;
}

export interface BackgroundTaskStatusSnapshot {
	taskId: string;
	role: string;
	agentId: string;
	status: 'running' | 'completed' | 'failed' | 'cancelled';
	startedAt: number;
	runId: string;
	messageGroupId?: string;
	threadId: string;
}

export interface RunStateTimeoutDetails {
	reason: InstanceAiLivenessTimeoutReason;
	surface: InstanceAiLivenessSurface;
	timeoutMs: number;
	elapsedMs: number;
	idleMs: number;
}

export interface StartRunOptions<TUser> {
	threadId: string;
	user: TUser;
	messageGroupId?: string;
}

export interface StartedRunState extends ActiveRunState {
	messageGroupId?: string;
}

export class RunStateRegistry<TUser = unknown> {
	private readonly activeRuns = new Map<string, ActiveRunState>();

	private readonly suspendedRuns = new Map<string, SuspendedRunState<TUser>>();

	private readonly pendingConfirmations = new Map<string, PendingConfirmation>();

	private readonly threadUsers = new Map<string, TUser>();

	private readonly threadMessageGroupId = new Map<string, string>();

	private readonly runIdsByMessageGroup = new Map<string, string[]>();

	/** IANA time zone captured at initial-run entry and reused by follow-up runs. */
	private readonly threadTimeZones = new Map<string, string>();

	/**
	 * Resolves a user id from the opaque `TUser` the registry is parameterised over.
	 * Required rather than optional: per-user concurrency counting depends on it, and a
	 * missing extractor would silently report every user as holding zero runs.
	 */
	constructor(private readonly getUserId: (user: TUser) => string) {}

	startRun(options: StartRunOptions<TUser>): StartedRunState {
		const runId = `run_${nanoid()}`;
		const abortController = new AbortController();
		const messageGroupId = options.messageGroupId ?? `mg_${nanoid()}`;
		const now = Date.now();

		this.activeRuns.set(options.threadId, {
			runId,
			threadId: options.threadId,
			abortController,
			messageGroupId,
			startedAt: now,
			lastActivityAt: now,
			userId: this.getUserId(options.user),
		});
		this.threadUsers.set(options.threadId, options.user);

		// When creating a fresh message group (no reuse), clean up the previous
		// one so runIdsByMessageGroup doesn't leak entries indefinitely.
		if (!options.messageGroupId) {
			const prevGroupId = this.threadMessageGroupId.get(options.threadId);
			if (prevGroupId && prevGroupId !== messageGroupId) {
				this.runIdsByMessageGroup.delete(prevGroupId);
			}
		}

		this.indexRunInGroup(options.threadId, messageGroupId, runId);

		return { runId, threadId: options.threadId, abortController, messageGroupId };
	}

	/**
	 * Seed the message-group indexes for a run: map the thread to its current
	 * group and record the run under that group. Idempotent.
	 *
	 * Called on `startRun` and re-applied on `suspendRun`/`activateSuspendedRun`
	 * so a run resumed after a restart (where these maps start empty) repopulates
	 * the group association the SSE bootstrap relies on.
	 */
	private indexRunInGroup(threadId: string, messageGroupId: string, runId: string): void {
		this.threadMessageGroupId.set(threadId, messageGroupId);
		let groupRunIds = this.runIdsByMessageGroup.get(messageGroupId);
		if (!groupRunIds) {
			groupRunIds = [];
			this.runIdsByMessageGroup.set(messageGroupId, groupRunIds);
		}
		if (!groupRunIds.includes(runId)) groupRunIds.push(runId);
	}

	getThreadStatus(
		threadId: string,
		backgroundTasks: BackgroundTaskStatusSnapshot[],
	): InstanceAiThreadStatusResponse {
		const activeRun = this.activeRuns.get(threadId);
		const suspendedRun = this.suspendedRuns.get(threadId);
		const liveRun = activeRun ?? suspendedRun;

		return {
			hasActiveRun: activeRun !== undefined,
			isSuspended: suspendedRun !== undefined,
			...(liveRun ? { runId: liveRun.runId } : {}),
			backgroundTasks: backgroundTasks
				.filter((task) => task.threadId === threadId)
				.map((task) => ({
					taskId: task.taskId,
					role: task.role,
					agentId: task.agentId,
					status: task.status,
					startedAt: task.startedAt,
					runId: task.runId,
					messageGroupId: task.messageGroupId,
				})),
		};
	}

	hasLiveRun(threadId: string): boolean {
		return this.activeRuns.has(threadId) || this.suspendedRuns.has(threadId);
	}

	hasActiveRun(threadId: string): boolean {
		return this.activeRuns.has(threadId);
	}

	hasSuspendedRun(threadId: string): boolean {
		return this.suspendedRuns.has(threadId);
	}

	getMessageGroupId(threadId: string): string | undefined {
		return this.threadMessageGroupId.get(threadId);
	}

	getLiveMessageGroupId(
		threadId: string,
		backgroundTasks: BackgroundTaskStatusSnapshot[],
	): string | undefined {
		if (this.hasLiveRun(threadId)) {
			return this.threadMessageGroupId.get(threadId);
		}

		const runningTask = backgroundTasks
			.filter((task) => task.threadId === threadId && task.status === 'running')
			.sort((a, b) => b.startedAt - a.startedAt)[0];

		return runningTask?.messageGroupId ?? this.threadMessageGroupId.get(threadId);
	}

	getRunIdsForMessageGroup(messageGroupId: string): string[] {
		return this.runIdsByMessageGroup.get(messageGroupId) ?? [];
	}

	getActiveRunId(threadId: string): string | undefined {
		return this.activeRuns.get(threadId)?.runId;
	}

	/**
	 * Runs holding a slot on this process, which is what the instance cap admits against.
	 *
	 * Includes a run parked on an inline approval card, which stays in `activeRuns`.
	 * Excludes a suspended run, which leaves `activeRuns` but keeps its agent in memory.
	 * So this bounds concurrent execution, not resident memory.
	 */
	activeRunCount(): number {
		return this.activeRuns.size;
	}

	/**
	 * Number of runs this user currently has executing, across all their threads.
	 *
	 * Excludes everything parked on a human, because none of it is spending anything and
	 * counting it would lock a user out for the whole confirmation timeout after a few
	 * abandoned cards. That means two things, not one:
	 *
	 *  - suspended runs, which have left `activeRuns` entirely; and
	 *  - runs blocked in `waitForConfirmation`, which stay in `activeRuns` while an inline
	 *    approval card is open. `sweepTimedOut` skips these for the same reason.
	 *
	 * {@link activeRunCount} differs on the second case only. An inline-parked run is still
	 * in `activeRuns`, so it holds a slot there. A suspended run holds a slot in neither,
	 * although it still retains its agent -- see {@link activeRunCount}.
	 *
	 * The scan is linear in concurrently-executing runs. That stays small in practice
	 * regardless of the caps -- each run costs ~12-20MB, so a process cannot hold many --
	 * and it is negligible next to the model call that follows.
	 */
	activeRunCountForUser(userId: string): number {
		let count = 0;
		for (const [threadId, run] of this.activeRuns) {
			if (run.userId !== userId) continue;
			if (this.hasPendingConfirmationForThread(threadId)) continue;
			count++;
		}
		return count;
	}

	getActiveRun(threadId: string): ActiveRunState | undefined {
		return this.activeRuns.get(threadId);
	}

	getSuspendedRun(threadId: string): SuspendedRunState<TUser> | undefined {
		return this.suspendedRuns.get(threadId);
	}

	attachTracing(threadId: string, tracing: InstanceAiTraceContext): void {
		const activeRun = this.activeRuns.get(threadId);
		if (!activeRun) return;

		this.activeRuns.set(threadId, {
			...activeRun,
			threadId,
			tracing,
		});
	}

	clearActiveRun(threadId: string, executionToken?: symbol): void {
		if (this.activeRuns.get(threadId)?.executionToken !== executionToken) return;
		this.activeRuns.delete(threadId);
	}

	cancelActiveRun(threadId: string): ActiveRunState | undefined {
		const active = this.activeRuns.get(threadId);
		if (!active) return undefined;

		this.activeRuns.delete(threadId);
		return active;
	}

	suspendRun(threadId: string, state: SuspendedRunState<TUser>): void {
		const activeRun = this.activeRuns.get(threadId);
		this.activeRuns.delete(threadId);
		state.startedAt = state.startedAt ?? activeRun?.startedAt ?? state.createdAt;
		state.lastActivityAt = state.lastActivityAt ?? state.createdAt;
		state.userId = state.userId ?? activeRun?.userId;
		this.suspendedRuns.set(threadId, state);

		// Re-seed group indexes: on a restart-resumed orphan these maps start
		// empty, so without this the SSE bootstrap loses the group association.
		if (state.messageGroupId) {
			this.indexRunInGroup(threadId, state.messageGroupId, state.runId);
		}
	}

	findSuspendedByRequestId(requestId: string): SuspendedRunState<TUser> | undefined {
		for (const run of this.suspendedRuns.values()) {
			if (run.requestId === requestId) return run;
		}
		return undefined;
	}

	cancelSuspendedRun(threadId: string): SuspendedRunState<TUser> | undefined {
		const suspended = this.suspendedRuns.get(threadId);
		if (!suspended) return undefined;

		this.suspendedRuns.delete(threadId);
		return suspended;
	}

	activateSuspendedRun(
		threadId: string,
		executionToken?: symbol,
	): SuspendedRunState<TUser> | undefined {
		const suspended = this.suspendedRuns.get(threadId);
		if (!suspended) return undefined;

		this.suspendedRuns.delete(threadId);
		const now = Date.now();
		this.activeRuns.set(threadId, {
			runId: suspended.runId,
			threadId,
			abortController: suspended.abortController,
			executionToken,
			messageGroupId: suspended.messageGroupId,
			tracing: suspended.tracing,
			modelId: suspended.modelId,
			startedAt: suspended.startedAt ?? suspended.createdAt,
			lastActivityAt: now,
			userId: suspended.userId ?? this.getUserId(suspended.user),
		});

		// Re-seed group indexes for the reactivated run (empty after a restart).
		if (suspended.messageGroupId) {
			this.indexRunInGroup(threadId, suspended.messageGroupId, suspended.runId);
		}
		return suspended;
	}

	registerPendingConfirmation(requestId: string, pending: PendingConfirmation): void {
		this.pendingConfirmations.set(requestId, {
			...pending,
			startedAt: pending.startedAt ?? pending.createdAt,
			lastActivityAt: pending.lastActivityAt ?? pending.createdAt,
		});
	}

	getPendingConfirmation(requestId: string): PendingConfirmation | undefined {
		return this.pendingConfirmations.get(requestId);
	}

	hasPendingConfirmationForThread(threadId: string): boolean {
		for (const pending of this.pendingConfirmations.values()) {
			if (pending.threadId === threadId) return true;
		}

		return false;
	}

	touchActiveRun(threadId: string, at = Date.now()): boolean {
		const active = this.activeRuns.get(threadId);
		if (!active) return false;
		active.lastActivityAt = at;
		return true;
	}

	touchSuspendedRun(threadId: string, at = Date.now()): boolean {
		const suspended = this.suspendedRuns.get(threadId);
		if (!suspended) return false;
		suspended.lastActivityAt = at;
		return true;
	}

	touchPendingConfirmation(requestId: string, at = Date.now()): boolean {
		const pending = this.pendingConfirmations.get(requestId);
		if (!pending) return false;
		pending.lastActivityAt = at;
		return true;
	}

	resolvePendingConfirmation(
		requestingUserId: string,
		requestId: string,
		data: ConfirmationData,
	): boolean {
		const pending = this.pendingConfirmations.get(requestId);
		if (!pending || pending.userId !== requestingUserId) return false;

		this.pendingConfirmations.delete(requestId);
		pending.resolve(data);
		return true;
	}

	cancelThread(
		threadId: string,
		cancelledConfirmation: ConfirmationData = { approved: false },
	): {
		active?: ActiveRunState;
		suspended?: SuspendedRunState<TUser>;
	} {
		for (const [requestId, pending] of this.pendingConfirmations) {
			if (pending.threadId !== threadId) continue;
			pending.resolve(cancelledConfirmation);
			this.pendingConfirmations.delete(requestId);
		}

		const active = this.activeRuns.get(threadId);
		const suspended = this.suspendedRuns.get(threadId);
		if (suspended) {
			this.suspendedRuns.delete(threadId);
		}

		return { ...(active ? { active } : {}), ...(suspended ? { suspended } : {}) };
	}

	getThreadUser(threadId: string): TUser | undefined {
		return this.threadUsers.get(threadId);
	}

	setTimeZone(threadId: string, timeZone: string): void {
		this.threadTimeZones.set(threadId, timeZone);
	}

	getTimeZone(threadId: string): string | undefined {
		return this.threadTimeZones.get(threadId);
	}

	/**
	 * Find active/suspended runs and pending confirmations that timed out under policy.
	 * Returns thread IDs and request IDs that should be cancelled/rejected.
	 * Does NOT mutate state — the caller is responsible for cancelling.
	 */
	sweepTimedOut(maxAgeMs: number): {
		suspendedThreadIds: string[];
		confirmationRequestIds: string[];
	};
	sweepTimedOut(
		policy: InstanceAiLivenessPolicy,
		now?: number,
	): {
		activeThreadIds: string[];
		suspendedThreadIds: string[];
		confirmationRequestIds: string[];
		activeTimeouts: Record<string, RunStateTimeoutDetails>;
		suspendedTimeouts: Record<string, RunStateTimeoutDetails>;
		confirmationTimeouts: Record<string, RunStateTimeoutDetails>;
	};
	sweepTimedOut(
		policyOrMaxAgeMs: InstanceAiLivenessPolicy | number,
		now = Date.now(),
	):
		| {
				suspendedThreadIds: string[];
				confirmationRequestIds: string[];
		  }
		| {
				activeThreadIds: string[];
				suspendedThreadIds: string[];
				confirmationRequestIds: string[];
				activeTimeouts: Record<string, RunStateTimeoutDetails>;
				suspendedTimeouts: Record<string, RunStateTimeoutDetails>;
				confirmationTimeouts: Record<string, RunStateTimeoutDetails>;
		  } {
		if (typeof policyOrMaxAgeMs === 'number') {
			const maxAgeMs = policyOrMaxAgeMs;
			const suspendedThreadIds: string[] = [];
			for (const [threadId, run] of this.suspendedRuns) {
				if (now - run.createdAt >= maxAgeMs) {
					suspendedThreadIds.push(threadId);
				}
			}
			const confirmationRequestIds: string[] = [];
			for (const [reqId, pending] of this.pendingConfirmations) {
				if (now - pending.createdAt >= maxAgeMs) {
					confirmationRequestIds.push(reqId);
				}
			}
			return { suspendedThreadIds, confirmationRequestIds };
		}

		const policy = policyOrMaxAgeMs;
		const activeThreadIds: string[] = [];
		const activeTimeouts: Record<string, RunStateTimeoutDetails> = {};
		for (const [threadId, run] of this.activeRuns) {
			if (this.hasPendingConfirmationForThread(threadId)) continue;

			const startedAt = run.startedAt ?? run.lastActivityAt ?? now;
			const lastActivityAt = run.lastActivityAt ?? startedAt;
			const decision = policy.evaluate({
				surface: 'active-run',
				startedAt,
				lastActivityAt,
				now,
			});
			if (decision.action === 'timeout') {
				activeThreadIds.push(threadId);
				activeTimeouts[threadId] = decision;
			}
		}

		const suspendedThreadIds: string[] = [];
		const suspendedTimeouts: Record<string, RunStateTimeoutDetails> = {};
		for (const [threadId, run] of this.suspendedRuns) {
			const startedAt = run.startedAt ?? run.createdAt;
			const lastActivityAt = run.lastActivityAt ?? run.createdAt;
			const decision = policy.evaluate({
				surface: 'suspended-run',
				startedAt,
				lastActivityAt,
				now,
			});
			if (decision.action === 'timeout') {
				suspendedThreadIds.push(threadId);
				suspendedTimeouts[threadId] = decision;
			}
		}
		const confirmationRequestIds: string[] = [];
		const confirmationTimeouts: Record<string, RunStateTimeoutDetails> = {};
		for (const [reqId, pending] of this.pendingConfirmations) {
			const startedAt = pending.startedAt ?? pending.createdAt;
			const lastActivityAt = pending.lastActivityAt ?? pending.createdAt;
			const decision = policy.evaluate({
				surface: 'pending-confirmation',
				startedAt,
				lastActivityAt,
				now,
			});
			if (decision.action === 'timeout') {
				confirmationRequestIds.push(reqId);
				confirmationTimeouts[reqId] = decision;
			}
		}
		return {
			activeThreadIds,
			suspendedThreadIds,
			confirmationRequestIds,
			activeTimeouts,
			suspendedTimeouts,
			confirmationTimeouts,
		};
	}

	/**
	 * Auto-reject a pending confirmation by request ID.
	 * Returns true if the confirmation existed and was rejected.
	 */
	rejectPendingConfirmation(requestId: string): boolean {
		const pending = this.pendingConfirmations.get(requestId);
		if (!pending) return false;
		this.pendingConfirmations.delete(requestId);
		pending.resolve({ approved: false });
		return true;
	}

	/** Remove a message-group entry from runIdsByMessageGroup. */
	deleteMessageGroup(groupId: string): void {
		this.runIdsByMessageGroup.delete(groupId);
	}

	/**
	 * Remove all per-thread state: active/suspended runs, confirmations,
	 * user, time zone, and message-group mappings.
	 * Returns the cancelled active/suspended runs so the caller can abort them.
	 */
	clearThread(
		threadId: string,
		cancelledConfirmation: ConfirmationData = { approved: false },
	): {
		active?: ActiveRunState;
		suspended?: SuspendedRunState<TUser>;
	} {
		const { active, suspended } = this.cancelThread(threadId, cancelledConfirmation);

		if (active) this.activeRuns.delete(threadId);

		this.threadUsers.delete(threadId);
		this.threadTimeZones.delete(threadId);

		const groupId = this.threadMessageGroupId.get(threadId);
		if (groupId) this.runIdsByMessageGroup.delete(groupId);
		this.threadMessageGroupId.delete(threadId);

		return { ...(active ? { active } : {}), ...(suspended ? { suspended } : {}) };
	}

	/**
	 * Process-wide teardown. Returns the in-flight runs so the service can
	 * abort them and persist terminal snapshots where appropriate.
	 *
	 * Pending confirmations are intentionally NOT resolved — auto-resolving an
	 * inline HITL (`waitForConfirmation(...)`) with `{ approved: false }`
	 * causes the awaiting agent tool to run to completion as "denied" before
	 * the process exits, which then mutates the snapshot tree mid-shutdown
	 * and clobbers the plan/ask card the user would otherwise see on reload.
	 * Letting the Promises dangle is safe: the process is exiting, the
	 * abortController for each active run is aborted next, and the
	 * `instance_ai_pending_confirmations` row survives so the user can still
	 * see the confirmation card (and get a clear "lost on restart" error if
	 * they click confirm — see `handleOrphanedConfirmation`).
	 *
	 * `pendingThreadIds` is returned so the service can skip the
	 * publish-run-finish + terminal-snapshot treatment for runs that are
	 * only sitting in `activeRuns` because they're waiting on an inline
	 * confirmation Promise.
	 */
	shutdown(): {
		activeRuns: ActiveRunState[];
		suspendedRuns: Array<SuspendedRunState<TUser>>;
		pendingThreadIds: string[];
	} {
		const activeRuns = [...this.activeRuns.values()];
		const suspendedRuns = [...this.suspendedRuns.values()];
		const pendingThreadIds = [
			...new Set([...this.pendingConfirmations.values()].map((p) => p.threadId)),
		];

		this.activeRuns.clear();
		this.suspendedRuns.clear();
		this.pendingConfirmations.clear();
		this.threadUsers.clear();
		this.threadTimeZones.clear();
		this.threadMessageGroupId.clear();
		this.runIdsByMessageGroup.clear();

		return { activeRuns, suspendedRuns, pendingThreadIds };
	}
}