import { resolve } from "node:path";
import {
	SessionManager,
	type ExtensionAPI,
	type SessionInfo,
} from "@earendil-works/pi-coding-agent";
import {
	Text,
	type AutocompleteItem,
	type AutocompleteProvider,
	type AutocompleteSuggestions,
	fuzzyFilter,
} from "@earendil-works/pi-tui";
import {
	SESSION_REFERENCE_CUSTOM_TYPE,
	SESSION_REFERENCE_PREFIX,
	buildReferenceContentFromSections,
	extractSessionReferenceIds,
	formatReferenceSession,
	sessionTitle,
	type ReferenceSessionInfo,
	type ReferenceSource,
} from "./session.ts";

const MAX_SESSION_SUGGESTIONS = 3;
const MAX_FILE_SUGGESTIONS = 7;
const MAX_REFERENCED_SESSIONS = 5;
const MENTION_PATTERN = /(?:^|[\t ])@([^\s@]*)$/;
const SUBAGENT_MANAGER_KEY = Symbol.for("pi-subagents:manager");

// ── In-process subagent record tracker ──────────────────────────────
// pi-subagents does not expose a global manager; we track records
// ourselves by listening to the events it emits.

type SubagentLiveRecord = {
	runId: string;
	sessionId: string;
	agent: string;
	cwd: string;
	startedAt: number;
	completedAt?: number;
};

type ReferenceDetails = {
	sessions: Array<{ id: string; title: string; cwd: string }>;
};

type SessionReference = {
	kind: "session" | "subagent";
	referenceIds: string[];
	info: ReferenceSessionInfo;
	path?: string;
	messages?: unknown[];
};

type SubagentRecord = {
	id: string;
	description?: string;
	startedAt?: number;
	completedAt?: number;
};

type SubagentManager = {
	getRecord(id: string): SubagentRecord | undefined;
};

// Local subagent record tracking (pi-subagents does not expose a global manager).
const liveSubagentRecords = new Map<string, SubagentLiveRecord>();
/** Cap retained completed/live records so long-lived pi processes cannot grow unbounded. */
const MAX_LIVE_SUBAGENT_RECORDS = 200;

function pruneLiveSubagentRecords(): void {
	while (liveSubagentRecords.size > MAX_LIVE_SUBAGENT_RECORDS) {
		const oldest = liveSubagentRecords.keys().next().value;
		if (oldest === undefined) break;
		liveSubagentRecords.delete(oldest);
	}
}

function clearLiveSubagentRecords(): void {
	liveSubagentRecords.clear();
}

function trackSubagentFromEvent(data: unknown): void {
	if (!data || typeof data !== "object") return;
	const event = data as Record<string, unknown>;
	// subagent:async-started payload has "id" as the run ID
	// subagent:async-complete payload has "runId" as the run ID
	const runId =
		(typeof event.id === "string" ? event.id : undefined) ??
		(typeof event.runId === "string" ? event.runId : undefined);
	if (!runId) return;
	const sessionId = typeof event.sessionId === "string" ? event.sessionId : undefined;
	if (!sessionId) return;

	const existing = liveSubagentRecords.get(runId);
	if (existing) {
		// Completion event — mark completedAt
		if (typeof event.endedAt === "number" || typeof event.lastUpdate === "number") {
			existing.completedAt = (
				typeof event.endedAt === "number" ? event.endedAt : event.lastUpdate
			) as number;
		}
		// Refresh insertion order so completed-but-still-referenced runs stay recent.
		liveSubagentRecords.delete(runId);
		liveSubagentRecords.set(runId, existing);
		return;
	}
	// Started event — create new record
	const agent = typeof event.agent === "string" ? event.agent : "";
	const cwd = typeof event.cwd === "string" ? event.cwd : "";
	const startedAt = typeof event.startedAt === "number" ? event.startedAt : Date.now();
	liveSubagentRecords.set(runId, {
		runId,
		sessionId,
		agent,
		cwd,
		startedAt,
	});
	pruneLiveSubagentRecords();
}

function getSubagentManager(): SubagentManager | undefined {
	// Try the global manager first (future-proof), fall back to local records.
	const manager = (globalThis as Record<PropertyKey, unknown>)[SUBAGENT_MANAGER_KEY] as
		| SubagentManager
		| undefined;
	if (manager && typeof manager.getRecord === "function") return manager;
	// If no global manager, use our local tracking.
	if (liveSubagentRecords.size === 0) return undefined;
	return {
		getRecord(id: string): SubagentRecord | undefined {
			const record = liveSubagentRecords.get(id);
			if (!record) return undefined;
			return {
				id: record.runId,
				description: record.agent || undefined,
				startedAt: record.startedAt,
				completedAt: record.completedAt,
			};
		},
	};
}

function extractMentionQuery(textBeforeCursor: string): string | undefined {
	return textBeforeCursor.match(MENTION_PATTERN)?.[1];
}

function formatDate(date: Date): string {
	return new Intl.DateTimeFormat(undefined, {
		year: "numeric",
		month: "2-digit",
		day: "2-digit",
	}).format(date);
}

/** Stable across keystrokes while getReferences returns the same reference objects. */
const sessionSearchTextCache = new WeakMap<SessionReference, string>();
const sessionItemCache = new WeakMap<
	SessionReference,
	{ cwd: string; named?: AutocompleteItem; stable?: AutocompleteItem }
>();

function sessionSearchText(reference: SessionReference): string {
	const cached = sessionSearchTextCache.get(reference);
	if (cached !== undefined) return cached;
	// 会话只按显式 session name 匹配，避免 firstMessage、路径或 ID 产生噪声。
	const text = reference.info.name?.trim() ?? "";
	sessionSearchTextCache.set(reference, text);
	return text;
}

function sessionItem(
	reference: SessionReference,
	currentCwd: string,
	useStableId: boolean,
): AutocompleteItem {
	let cached = sessionItemCache.get(reference);
	const variant = useStableId ? "stable" : "named";
	if (cached?.cwd === currentCwd && cached[variant]) return cached[variant];

	const session = reference.info;
	const workspace = samePath(session.cwd, currentCwd)
		? "current workspace"
		: session.cwd || "unknown workspace";
	const label = reference.kind === "subagent" ? "[SubAgent]" : "[Session]";
	// 唯一名称保持短格式；同名或无名称时使用稳定 ID，避免引用歧义。
	const sessionName = session.name?.trim();
	const referenceId = useStableId
		? reference.referenceIds[0]
		: (sessionName ?? reference.referenceIds[0]);
	const item: AutocompleteItem = {
		value: `${SESSION_REFERENCE_PREFIX}[${referenceId}]`,
		label: `${label} ${sessionTitle(session)}`,
		description: `${workspace} · ${session.messageCount} messages · ${formatDate(session.modified)}`,
	};
	if (cached?.cwd !== currentCwd) cached = { cwd: currentCwd };
	cached[variant] = item;
	sessionItemCache.set(reference, cached);
	return item;
}

/** Sort cache: same array + cwd reuses order (getReferences returns a stable ordered array). */
const orderedReferencesCache = new WeakMap<
	SessionReference[],
	{ cwd: string; ordered: SessionReference[] }
>();

function orderSessionReferences(
	references: SessionReference[],
	currentCwd: string,
): SessionReference[] {
	const cached = orderedReferencesCache.get(references);
	if (cached && cached.cwd === currentCwd) return cached.ordered;

	const ordered = [...references].sort((left, right) => {
		const leftLocal = samePath(left.info.cwd, currentCwd) ? 1 : 0;
		const rightLocal = samePath(right.info.cwd, currentCwd) ? 1 : 0;
		return rightLocal - leftLocal || right.info.modified.getTime() - left.info.modified.getTime();
	});
	orderedReferencesCache.set(references, { cwd: currentCwd, ordered });
	// Point the ordered array at itself so subsequent keystrokes skip re-sorting.
	orderedReferencesCache.set(ordered, { cwd: currentCwd, ordered });
	return ordered;
}

function filterSessions(
	references: SessionReference[],
	query: string,
	currentCwd: string,
): AutocompleteItem[] {
	if (query.startsWith("session:")) return [];

	const ordered = orderSessionReferences(references, currentCwd);
	const nameCounts = new Map<string, number>();
	for (const reference of ordered) {
		const name = reference.info.name?.trim();
		if (name) nameCounts.set(name, (nameCounts.get(name) ?? 0) + 1);
	}
	const trimmed = query.trim();
	// 模糊匹配只接受有 session name 的会话；subagent 按 agent name 匹配。
	const searchable = ordered.filter(
		(reference) => reference.kind === "subagent" || Boolean(reference.info.name?.trim()),
	);
	// Empty query: top-N already sorted — no fuzzy scan over the full list。
	const matches = trimmed
		? fuzzyFilter(searchable, trimmed, sessionSearchText)
		: ordered.slice(0, MAX_SESSION_SUGGESTIONS);
	return matches.slice(0, MAX_SESSION_SUGGESTIONS).map((reference) => {
		const name = reference.info.name?.trim();
		return sessionItem(reference, currentCwd, Boolean(name && (nameCounts.get(name) ?? 0) > 1));
	});
}

function isPathLikeQuery(query: string): boolean {
	return /[\\/*?]/.test(query) || /\.[A-Za-z0-9_-]{1,12}$/.test(query);
}

function mergeSessionAndFileItems(
	sessionItems: AutocompleteItem[],
	fileItems: AutocompleteItem[],
	query: string,
): AutocompleteItem[] {
	const sessions = sessionItems.slice(0, MAX_SESSION_SUGGESTIONS);
	const files = fileItems.slice(0, MAX_FILE_SUGGESTIONS);
	if (isPathLikeQuery(query)) return [...files, ...sessions];

	const merged: AutocompleteItem[] = [];
	let sessionIndex = 0;
	let fileIndex = 0;
	while (sessionIndex < sessions.length || fileIndex < files.length) {
		for (let count = 0; count < 2 && fileIndex < files.length; count++) {
			merged.push(files[fileIndex++]!);
		}
		if (sessionIndex < sessions.length) merged.push(sessions[sessionIndex++]!);
	}
	return merged;
}

function liveSubagentReferences(
	agentIds: Set<string>,
	currentSessionId: string,
): SessionReference[] {
	const manager = getSubagentManager();
	if (!manager) return [];

	const references: SessionReference[] = [];
	for (const agentId of agentIds) {
		const record = manager.getRecord(agentId);
		if (!record) continue;

		// Try the global manager's live session first (future-proof).
		const liveRecord = liveSubagentRecords.get(agentId);
		const sessionId = liveRecord?.sessionId ?? record.id;
		if (!sessionId || sessionId === currentSessionId) continue;

		const name = record.description?.trim() || liveRecord?.agent || undefined;
		const modifiedAt = record.completedAt ?? record.startedAt ?? Date.now();
		references.push({
			kind: "subagent",
			referenceIds: [sessionId, agentId],
			info: {
				id: sessionId,
				name,
				cwd: liveRecord?.cwd ?? "",
				firstMessage: "",
				messageCount: 0,
				modified: new Date(modifiedAt),
			},
		});
	}
	return references;
}

function mergeReferences(
	sessions: SessionInfo[],
	subagents: SessionReference[],
): SessionReference[] {
	const bySessionId = new Map<string, SessionReference>();
	for (const session of sessions) {
		bySessionId.set(session.id, {
			kind: "session",
			referenceIds: [session.id],
			info: session,
			path: session.path,
		});
	}
	for (const subagent of subagents) {
		// Prefer the live subagent context when a persisted subagent has the same ID.
		bySessionId.set(subagent.info.id, subagent);
	}
	return [...bySessionId.values()];
}

export function createAutocompleteProvider(
	current: AutocompleteProvider,
	getReferences: () => Promise<SessionReference[]>,
	currentCwd: string,
): AutocompleteProvider {
	return {
		triggerCharacters: ["@"],
		async getSuggestions(
			lines,
			cursorLine,
			cursorCol,
			options,
		): Promise<AutocompleteSuggestions | null> {
			const currentLine = lines[cursorLine] ?? "";
			const query = extractMentionQuery(currentLine.slice(0, cursorCol));
			if (query === undefined) {
				return current.getSuggestions(lines, cursorLine, cursorCol, options);
			}

			const [baseSuggestions, references] = await Promise.all([
				current.getSuggestions(lines, cursorLine, cursorCol, options),
				getReferences(),
			]);
			if (options.signal.aborted) return null;

			const sessionItems = filterSessions(references, query, currentCwd);
			const fileItems = baseSuggestions?.prefix === `@${query}` ? baseSuggestions.items : [];
			const items = mergeSessionAndFileItems(sessionItems, fileItems, query);
			if (items.length === 0) return baseSuggestions;
			return { prefix: `@${query}`, items };
		},

		applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
			return current.applyCompletion(lines, cursorLine, cursorCol, item, prefix);
		},

		shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
			return current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ?? true;
		},
	};
}

function samePath(left: string | undefined, right: string): boolean {
	if (!left) return false;
	const normalizedLeft = resolve(left);
	const normalizedRight = resolve(right);
	return process.platform === "win32"
		? normalizedLeft.toLowerCase() === normalizedRight.toLowerCase()
		: normalizedLeft === normalizedRight;
}

export default function sessionReferenceExtension(pi: ExtensionAPI): void {
	let getAvailableReferences: (() => Promise<SessionReference[]>) | undefined;
	let sessionGeneration = 0;
	const subagentIds = new Set<string>();
	// pi-subagents emits "subagent:async-started" and "subagent:async-complete" events.
	// We track records locally so @[SubAgent] suggestions work even without a global manager.
	const subagentEventNames = ["subagent:async-started", "subagent:async-complete"];
	const unsubscribeSubagentEvents = subagentEventNames.map((eventName) =>
		pi.events.on(eventName, (data) => {
			trackSubagentFromEvent(data);
			if (!data || typeof data !== "object") return;
			const payload = data as Record<string, unknown>;
			// subagent:async-started uses "id", subagent:async-complete uses "runId"
			const id =
				(typeof payload.id === "string" ? payload.id : undefined) ??
				(typeof payload.runId === "string" ? payload.runId : undefined);
			if (typeof id === "string" && id) subagentIds.add(id);
		}),
	);

	pi.registerMessageRenderer(SESSION_REFERENCE_CUSTOM_TYPE, (message, _options, theme) => {
		const details = message.details as ReferenceDetails | undefined;
		const sessions = details?.sessions ?? [];
		const labels = sessions.map((session) => session.title).join(", ");
		const summary =
			sessions.length === 1 ? "Referenced 1 session" : `Referenced ${sessions.length} sessions`;
		const text = labels ? `${theme.fg("accent", summary)}\n${theme.fg("dim", labels)}` : summary;
		return new Text(text, 1, 0);
	});

	pi.on("session_start", (_event, ctx) => {
		const generation = ++sessionGeneration;
		// 会话替换/reload 会让 ctx 的 getter 抛 stale 错误，await 之后不能再读；
		// 这里在同步阶段一次性取出纯值。
		const currentCwd = ctx.cwd;
		const ui = ctx.ui;
		subagentIds.clear();
		clearLiveSubagentRecords();
		let loadErrorShown = false;
		const currentSessionId = ctx.sessionManager.getSessionId();
		const currentSessionFile = ctx.sessionManager.getSessionFile();
		let sessionsPromise: Promise<SessionInfo[]> | undefined;

		const getSessions = (): Promise<SessionInfo[]> => {
			sessionsPromise ||= SessionManager.listAll()
				.then((sessions) =>
					sessions.filter(
						(session) =>
							session.id !== currentSessionId && !samePath(currentSessionFile, session.path),
					),
				)
				.catch((error: unknown) => {
					if (!loadErrorShown) {
						loadErrorShown = true;
						const reason = error instanceof Error ? error.message : String(error);
						ui.notify(`session-reference: failed to load sessions: ${reason}`, "error");
					}
					return [];
				});
			return sessionsPromise;
		};

		// Pre-sort once per (sessions list, subagent id set) so @ keystrokes only fuzzy-filter.
		let referencesCache:
			| { sessions: SessionInfo[]; subagentKey: string; ordered: SessionReference[] }
			| undefined;
		const getReferences = async (): Promise<SessionReference[]> => {
			const sessions = await getSessions();
			// 会话替换/reload 后丢弃旧 generation 的结果，避免对失效状态继续工作。
			if (generation !== sessionGeneration) return [];
			const subagentKey = [...subagentIds].join("\0");
			if (
				referencesCache &&
				referencesCache.sessions === sessions &&
				referencesCache.subagentKey === subagentKey
			) {
				return referencesCache.ordered;
			}
			const ordered = orderSessionReferences(
				mergeReferences(sessions, liveSubagentReferences(subagentIds, currentSessionId)),
				currentCwd,
			);
			referencesCache = { sessions, subagentKey, ordered };
			return ordered;
		};

		getAvailableReferences = getReferences;
		if (ctx.mode === "tui") {
			// 预取是 detached 的，必须兜住 rejection，否则 stale 错误会变成
			// unhandled rejection 直接终止 Pi。
			void getReferences().catch(() => {});
			// Register after other session_start handlers. pi-fff claims every @
			// prefix, so a provider installed before it would never see session mentions.
			setTimeout(() => {
				if (generation !== sessionGeneration) return;
				ui.addAutocompleteProvider((current) =>
					createAutocompleteProvider(current, getReferences, currentCwd),
				);
			}, 0);
		}
	});

	pi.on("before_agent_start", async (event, ctx) => {
		const referenceIds = extractSessionReferenceIds(event.prompt);
		if (referenceIds.length === 0) return;

		const currentSessionId = ctx.sessionManager.getSessionId();
		const references = await (getAvailableReferences?.() ??
			SessionManager.listAll().then((sessions) =>
				mergeReferences(
					sessions.filter((session) => session.id !== currentSessionId),
					liveSubagentReferences(subagentIds, currentSessionId),
				),
			));
		const referencesById = new Map<string, SessionReference>();
		const referencesByName = new Map<string, SessionReference>();
		const ambiguousNames = new Set<string>();
		for (const reference of references) {
			for (const id of reference.referenceIds) referencesById.set(id, reference);
			const name = reference.info.name?.trim();
			if (!name || ambiguousNames.has(name)) continue;
			if (referencesByName.has(name)) {
				referencesByName.delete(name);
				ambiguousNames.add(name);
			} else {
				referencesByName.set(name, reference);
			}
		}
		const seenReferences = new Set<SessionReference>();
		const matchingReferences = referenceIds
			.map((id) => referencesById.get(id) ?? referencesByName.get(id))
			.filter((reference): reference is SessionReference => {
				if (!reference || reference.info.id === currentSessionId || seenReferences.has(reference))
					return false;
				seenReferences.add(reference);
				return true;
			});
		const selected = matchingReferences.slice(0, MAX_REFERENCED_SESSIONS);

		if (selected.length === 0) {
			ctx.ui.notify("session-reference: referenced sessions were not found", "warning");
			return;
		}

		if (matchingReferences.length > MAX_REFERENCED_SESSIONS) {
			ctx.ui.notify(
				`session-reference: only the first ${MAX_REFERENCED_SESSIONS} sessions were included`,
				"warning",
			);
		}

		const sections: string[] = [];
		const referencedSessions: ReferenceDetails["sessions"] = [];
		for (const reference of selected) {
			try {
				const info = reference.info;
				const messages =
					reference.messages ??
					(reference.path
						? SessionManager.open(reference.path).buildSessionContext().messages
						: (() => {
								throw new Error("reference session is no longer available");
							})());
				const source: ReferenceSource = { info, messages };
				sections.push(formatReferenceSession(source));
				referencedSessions.push({
					id: info.id,
					title: sessionTitle(info),
					cwd: info.cwd,
				});
			} catch (error) {
				const reason = error instanceof Error ? error.message : String(error);
				ctx.ui.notify(
					`session-reference: failed to read ${reference.info.id}: ${reason}`,
					"warning",
				);
			}
		}
		if (sections.length === 0) return;

		const details: ReferenceDetails = { sessions: referencedSessions };
		return {
			message: {
				customType: SESSION_REFERENCE_CUSTOM_TYPE,
				content: buildReferenceContentFromSections(sections),
				display: true,
				details,
			},
		};
	});

	pi.on("session_before_switch", () => {
		subagentIds.clear();
		clearLiveSubagentRecords();
		getAvailableReferences = undefined;
	});

	pi.on("session_shutdown", () => {
		sessionGeneration++;
		for (const unsubscribe of unsubscribeSubagentEvents) unsubscribe();
		subagentIds.clear();
		clearLiveSubagentRecords();
		getAvailableReferences = undefined;
	});
}