import type { EditorFactory } from "./pi-adapter.ts";
import type { ImageContent } from "./content.ts";
import {
	createImageMarkerLink,
	renderImageMarkerLinks,
	sanitizeModelMessages,
} from "./attachment-links.ts";
import { type GalleryImage } from "./image-gallery.ts";
import { GalleryPresenter } from "./gallery-presenter.ts";
import { extractImagePaths } from "./image-paths.ts";
import { upgradeScreenshotToolResult } from "./tool-result-upgrader.ts";
import { debugLog } from "./debug.ts";
import { PreparedImage } from "./image-pipeline.ts";
import { DraftStore, type DraftImage } from "./draft-store.ts";
import { PiEditorAdapter } from "./pi-adapter.ts";
import { supportsInlinePreview } from "./image-gallery.ts";

// ── Types ──────────────────────────────────────────────────

type TrackedImage = DraftImage;

export type ExtensionDeps = {
	readImageContentFromPathAsync: (
		filePath: string,
	) => Promise<ImageContent | null>;
	/** Terminal thumbnail only; never used as model input. */
	maybeResizeImage?: (image: ImageContent) => Promise<ImageContent>;
	resizeModelImage?: (image: ImageContent) => Promise<ImageContent>;
	normalizeImageForMatching?: (image: ImageContent) => Promise<ImageContent>;
	createAtomicEditor?: (tui: Parameters<EditorFactory>[0], theme: Parameters<EditorFactory>[1], keybindings: Parameters<EditorFactory>[2], attachImage: (image: ImageContent) => string, baseEditor?: ReturnType<EditorFactory>, draftReference?: (marker: string) => string | undefined) => ReturnType<EditorFactory>;
	/** Resolves before installing the custom editor so unavailable clipboard backends keep Pi paste intact. */
	supportsAtomicEditor?: () => Promise<boolean>;
	isImagePasteInput?: (data: string) => boolean;
	loadImageContentFromPath: (
		filePath: string,
	) => Promise<ImageContent | null>;
	storeImage?: (image: ImageContent) => Promise<string | undefined>;
	resolveImageReference?: (reference: string) => string | undefined;
};

type PiLike = {
	on(event: string, handler: (...args: any[]) => any): void;
	registerCommand?: (name: string, options: { description?: string; handler: (args: string, ctx: CtxLike) => Promise<void> }) => void;
	registerMarkdownTransformer?: (
		transformer: (markdown: string, context: { messageType: string }) => string,
	) => void;
};

type CtxLike = {
	cwd: string;
	hasUI?: boolean;
	mode?: "tui" | "rpc" | "json" | "print";
	sessionManager?: { getBranch(): readonly unknown[] };
	isIdle(): boolean;
	ui: {
		getEditorComponent?(): EditorFactory | undefined;
		onTerminalInput?(handler: (data: string) => { consume?: boolean; data?: string } | undefined): () => void;
		setEditorComponent?(factory: EditorFactory | undefined): void;
		setWidget(
			key: string,
			content:
				| string[]
				| ((tui: any, theme: any) => any)
				| undefined,
			options?: { placement?: "aboveEditor" | "belowEditor" },
		): void;
		getEditorText(): string;
		setEditorText(text: string): void;
		notify(message: string, type?: "info" | "warning" | "error"): void;
		theme: any;
	};
};

/** Event shape for the "input" event from pi. */
type InputEvent = {
	text: string;
	images?: ImageContent[];
};

/** Discriminated union for input handler return values. */
type InputResult =
	| { action: "continue" }
	| { action: "handled" }
	| { action: "transform"; text: string; images: ImageContent[] };

/** Re-export for tool_result event typing. */
type ToolResultEvent = import("./tool-result-upgrader.ts").ToolResultEventLike;

// ── Constants ──────────────────────────────────────────────

const WIDGET_KEY = "image-view";
const POLL_INTERVAL_MS = 250;
const IMAGE_PLACEHOLDER_RE = /\[Image #\d+\]/g;

function placeholdersIn(text: string): Set<string> {
	return new Set(text.match(IMAGE_PLACEHOLDER_RE) ?? []);
}

function highestImageNumberInText(text: string): number {
	let highest = 0;
	for (const match of text.matchAll(/\[Image #(\d+)\]/g)) {
		const value = Number(match[1]);
		if (Number.isSafeInteger(value)) highest = Math.max(highest, value);
	}
	return highest;
}

function highestImageNumberInContent(content: unknown): number {
	if (typeof content === "string") return highestImageNumberInText(content);
	if (!Array.isArray(content)) return 0;
	return content.reduce((highest, part) => {
		if (typeof part === "string") return Math.max(highest, highestImageNumberInText(part));
		if (!part || typeof part !== "object") return highest;
		const text = (part as { text?: unknown }).text;
		return typeof text === "string"
			? Math.max(highest, highestImageNumberInText(text))
			: highest;
	}, 0);
}

export function nextImageNumberForBranch(entries: readonly unknown[]): number {
	let highest = 0;
	for (const entry of entries) {
		if (!entry || typeof entry !== "object") continue;
		const candidate = entry as { type?: unknown; message?: { role?: unknown; content?: unknown } };
		if (candidate.type !== "message" || candidate.message?.role !== "user") continue;
		highest = Math.max(highest, highestImageNumberInContent(candidate.message.content));
	}
	return highest + 1;
}

// ── Extension ──────────────────────────────────────────────

export function registerImagePreviewExtension(
	pi: PiLike,
	deps: ExtensionDeps,
): void {
	const resolveImageReference = deps.resolveImageReference;
	if (pi.registerMarkdownTransformer && resolveImageReference) {
		pi.registerMarkdownTransformer((markdown, context) =>
			context.messageType === "user"
				? renderImageMarkerLinks(markdown, resolveImageReference)
				: markdown,
		);
	}

	let clearBeforeIndex: number | undefined;
	let clearOnNextContext = false;
	let lastContextMessageCount = 0;

	pi.registerCommand?.("pi-image-view", {
		description: "Clear existing image context (detail is a 2000px compatibility alias)",
		handler: async (args, ctx) => {
			const action = args.trim();
			if (action === "detail") {
				ctx.ui.notify("Images already use up to 2000px by default", "info");
				return;
			}
			if (action !== "clear") {
				ctx.ui.notify("Usage: /pi-image-view [clear|detail]", "error");
				return;
			}
			if (lastContextMessageCount > 0) {
				clearBeforeIndex = lastContextMessageCount;
				clearOnNextContext = false;
			} else {
				clearOnNextContext = true;
			}
			ctx.ui.notify("Existing images cleared from model context", "info");
		},
	});

	pi.on("context", (event: { messages: Array<{ role?: unknown; content?: unknown }> }) => {
		if (clearBeforeIndex !== undefined && event.messages.length < lastContextMessageCount) {
			const removedPrefixLength = lastContextMessageCount - event.messages.length;
			clearBeforeIndex = Math.max(0, clearBeforeIndex - removedPrefixLength);
		}
		if (clearOnNextContext) {
			let latestUserIndex = -1;
			for (let index = event.messages.length - 1; index >= 0; index -= 1) {
				if (event.messages[index]?.role === "user") {
					latestUserIndex = index;
					break;
				}
			}
			clearBeforeIndex = latestUserIndex >= 0 ? latestUserIndex : event.messages.length;
			clearOnNextContext = false;
		}
		lastContextMessageCount = event.messages.length;
		return { messages: sanitizeModelMessages(event.messages, clearBeforeIndex) };
	});

	const editorAdapter = new PiEditorAdapter();
	const draft = new DraftStore();
	let presenter: GalleryPresenter | null = null;
	let pollTimer: ReturnType<typeof setInterval> | null = null;
	let unsubscribeTerminalInput: (() => void) | undefined;
	const pasteScanTimers = new Set<ReturnType<typeof setTimeout>>();
	let latestCtx: CtxLike | null = null;
	let scanInFlight = false;
	let scanGeneration = 0;
	let lastScannedText: string | undefined;
	let failedScanText: string | undefined;
	let failedScanAttempts = 0;
	let sessionGeneration = 0;

	// ── Helpers ────────────────────────────────────────────


	function allocatePlaceholder(existingText = "", existing = draft.images): string {
		return draft.allocate(existingText, existing);
	}

	function prepare(image: ImageContent): PreparedImage {
		return new PreparedImage(image, {
			resizeModelImage: deps.resizeModelImage,
			resizePreviewImage: deps.maybeResizeImage,
			storeImage: deps.storeImage,
		});
	}

	function startPreparation(entry: TrackedImage): void {
		void entry.prepared.persist().then(() => {
			if (latestCtx && draft.has(entry)) editorAdapter.refresh();
		});
		if (supportsInlinePreview()) {
			void entry.prepared.thumbnail().then(() => {
				if (latestCtx && draft.has(entry)) refreshWidget(latestCtx);
			});
		}
	}

	function attachClipboardImage(image: ImageContent): string {
		let editorText = "";
		try { editorText = latestCtx?.ui.getEditorText() ?? ""; } catch { /* editor unavailable during teardown */ }
		const placeholder = allocatePlaceholder(editorText);
		const entry: TrackedImage = { placeholder, prepared: prepare(image) };
		draft.images.set(placeholder, entry);
		startPreparation(entry);
		return placeholder;
	}

	function refreshWidget(ctx: CtxLike): void {
		if (ctx.hasUI === false || (ctx.mode !== undefined && ctx.mode !== "tui")) return;
		const galleryImages: GalleryImage[] = [];
		if (supportsInlinePreview()) {
			for (const entry of draft.images.values()) {
				const preview = entry.prepared.preview;
				if (preview?.mimeType === "image/png") galleryImages.push({
					key: entry.placeholder, data: preview.data, mimeType: preview.mimeType,
					label: entry.placeholder,
				});
			}
		}
		if (galleryImages.length === 0 && !presenter) return;
		if (!presenter) {
			presenter = new GalleryPresenter({ widget: ctx.ui, key: WIDGET_KEY });
		}
		presenter.update(galleryImages);
	}

	function resetDraft(ctx: CtxLike): void {
		presenter?.dispose();
		presenter = null;
		draft.reset();
		scanGeneration += 1;
		lastScannedText = undefined;
		failedScanText = undefined;
		failedScanAttempts = 0;
		if (ctx.hasUI !== false && (ctx.mode === undefined || ctx.mode === "tui")) {
			ctx.ui.setWidget(WIDGET_KEY, undefined);
		}
	}

	/**
	 * Scan editor text for image paths.
	 * Track new ones, remove ones that are no longer in the text.
	 * Async to avoid blocking the event loop with file I/O.
	 */
	async function scanEditorText(ctx: CtxLike): Promise<void> {
		let text: string;
		try {
			text = ctx.ui.getEditorText();
		} catch (err) {
			debugLog("Failed to get editor text", err);
			return;
		}
		if (text === lastScannedText) return;
		if (text !== failedScanText) {
			failedScanText = text;
			failedScanAttempts = 0;
		}
		lastScannedText = text;
		const generation = ++scanGeneration;
		if (scanInFlight) return;

		scanInFlight = true;
		try {
			const visiblePlaceholders = placeholdersIn(text);
			const nextTracked = new Map(draft.images);
			let changed = false;
			for (const placeholder of nextTracked.keys()) {
				if (!visiblePlaceholders.has(placeholder)) {
					nextTracked.delete(placeholder);
					changed = true;
				}
			}

			let renderedText = text;
			let hadReadFailure = false;
			const newEntries: TrackedImage[] = [];
			for (const { raw, path: filePath } of extractImagePaths(text, { cwd: ctx.cwd })) {
				const image = await deps.readImageContentFromPathAsync(filePath);
				if (generation !== scanGeneration) {
					lastScannedText = undefined;
					return;
				}
				if (!image) {
					hadReadFailure = true;
					continue;
				}

				const placeholder = allocatePlaceholder(renderedText, nextTracked);

				const entry: TrackedImage = { placeholder, prepared: prepare(image) };
				nextTracked.set(placeholder, entry);
				newEntries.push(entry);
				renderedText = renderedText.replace(raw, placeholder);
				changed = true;
			}

			if (generation !== scanGeneration || ctx.ui.getEditorText() !== text) {
				lastScannedText = undefined;
				return;
			}
			draft.images = nextTracked;
			if (renderedText !== text) {
				ctx.ui.setEditorText(renderedText);
				lastScannedText = renderedText;
			}
			if (changed) refreshWidget(ctx);
			for (const entry of newEntries) startPreparation(entry);
			if (hadReadFailure) {
				failedScanAttempts += 1;
				if (failedScanAttempts < 3) lastScannedText = undefined;
			} else {
				failedScanText = undefined;
				failedScanAttempts = 0;
			}
		} finally {
			scanInFlight = false;
		}
	}

	function schedulePasteScans(ctx: CtxLike): void {
		for (const timer of pasteScanTimers) clearTimeout(timer);
		pasteScanTimers.clear();
		for (const delay of [0, 16, 40, 80, 160, 240]) {
			const timer = setTimeout(() => {
				pasteScanTimers.delete(timer);
				scanEditorText(ctx).catch((error) => debugLog("Paste-triggered image scan failed", error));
			}, delay);
			timer.unref?.();
			pasteScanTimers.add(timer);
		}
	}

	function startPolling(): void {
		stopPolling();
		pollTimer = setInterval(() => {
			if (!latestCtx) return;
			scanEditorText(latestCtx).catch((err) => {
				debugLog("Error during editor text scan", err);
			});
		}, POLL_INTERVAL_MS);
		// Don't let the poll timer keep the Node event loop alive.
		// In non-interactive modes (e.g. `pi --print`), the process must be
		// able to exit once work is done — an active interval would block exit.
		if (typeof pollTimer.unref === "function") {
			pollTimer.unref();
		}
	}

	function stopPolling(): void {
		if (pollTimer) {
			clearInterval(pollTimer);
			pollTimer = null;
		}
	}

	// ── Event handlers ─────────────────────────────────────

	const cleanup = (): void => {
		stopPolling();
		unsubscribeTerminalInput?.();
		unsubscribeTerminalInput = undefined;
		for (const timer of pasteScanTimers) clearTimeout(timer);
		pasteScanTimers.clear();
		scanGeneration += 1;
		sessionGeneration += 1;
		lastScannedText = undefined;
		latestCtx = null;
		draft.reset();
		presenter?.dispose();
		presenter = null;
	};

	pi.on("session_start", async (_event: unknown, ctx: CtxLike) => {
		const generation = ++sessionGeneration;
		stopPolling();
		unsubscribeTerminalInput?.();
		unsubscribeTerminalInput = undefined;
		for (const timer of pasteScanTimers) clearTimeout(timer);
		pasteScanTimers.clear();
		clearBeforeIndex = undefined;
		clearOnNextContext = false;
		lastContextMessageCount = 0;
		latestCtx = ctx;
		resetDraft(ctx);
		draft.reset(nextImageNumberForBranch(ctx.sessionManager?.getBranch() ?? []));
		const interactive = ctx.hasUI !== false && (ctx.mode === undefined || ctx.mode === "tui");
		const createAtomicEditor = deps.createAtomicEditor;
		if (interactive && createAtomicEditor && ctx.ui.setEditorComponent) {
			let supported = true;
			try {
				if (deps.supportsAtomicEditor) supported = await deps.supportsAtomicEditor();
			} catch (error) {
				supported = false;
				debugLog("Direct clipboard capability detection failed", error);
			}
			// A shutdown or a newer session may have won while capability detection ran.
			// Stop the obsolete handler before *all* later startup side effects.
			if (generation !== sessionGeneration || latestCtx !== ctx) return;
			if (supported) {
				editorAdapter.install(ctx.ui, (tui, theme, keys, base) =>
					createAtomicEditor(tui, theme, keys, attachClipboardImage, base, (marker) => draft.reference(marker)));
			}
		}
		if (interactive) {
			startPolling();
			const isImagePasteInput = deps.isImagePasteInput;
			if (isImagePasteInput && ctx.ui.onTerminalInput) {
				unsubscribeTerminalInput = ctx.ui.onTerminalInput((data) => {
					if (isImagePasteInput(data)) schedulePasteScans(ctx);
					return undefined;
				});
			}
		}
	});

	pi.on("session_shutdown", (_event: unknown, ctx: CtxLike) => {
		cleanup();
		editorAdapter.dispose(ctx.ui);
	});

	pi.on("tool_result", async (event: ToolResultEvent, ctx: CtxLike) => {
		return upgradeScreenshotToolResult(
			event,
			ctx.cwd,
			deps.loadImageContentFromPath,
			deps.resizeModelImage,
		);
	});

	// On submit: remove local placeholders/paths and attach the same image content.
	pi.on("input", async (event: InputEvent, ctx: CtxLike): Promise<InputResult> => {
		const inputSession = sessionGeneration;
		const fullText = (event.text || "").trim();

		const detectedPaths = extractImagePaths(fullText, { cwd: ctx.cwd });
		if (
			(fullText.startsWith("/") && detectedPaths.length === 0) ||
			fullText.trimStart().startsWith("!")
		) {
			return { action: "continue" };
		}

		const candidates: Array<{ token: string; entry: TrackedImage; index: number }> = [];
		for (const [placeholder, entry] of draft.images) {
			const index = fullText.indexOf(placeholder);
			if (index >= 0) candidates.push({ token: placeholder, entry, index });
		}

		// Fast-submit fallback: the input event may arrive before the 250ms editor
		// poll has converted a freshly pasted path into a placeholder.
		for (const { raw, path: filePath } of detectedPaths) {
			const image = await deps.readImageContentFromPathAsync(filePath);
			if (!image) continue;
			const placeholder = allocatePlaceholder(fullText);
			candidates.push({
				token: raw,
				index: fullText.indexOf(raw),
				entry: { placeholder, prepared: prepare(image) },
			});
		}

		if (candidates.length === 0) return { action: "continue" };
		candidates.sort((a, b) => a.index - b.index);

		const existingByContent = new Map<string, number[]>();
		for (const [index, image] of (event.images ?? []).entries()) {
			const key = `${image.mimeType}\u0000${image.data}`;
			const matches = existingByContent.get(key);
			if (matches) matches.push(index);
			else existingByContent.set(key, [index]);
		}
		const existingIndexes = candidates.map(({ entry }) => {
			const source = entry.prepared.source;
			const key = `${source.mimeType}\u0000${source.data}`;
			return existingByContent.get(key)?.shift();
		});
		const normalizeImageForMatching = deps.normalizeImageForMatching;
		if (normalizeImageForMatching && (event.images?.length ?? 0) > 0) {
			const normalizedCandidates = await Promise.all(
				candidates.map(({ entry }, candidateIndex) =>
					existingIndexes[candidateIndex] === undefined
						? normalizeImageForMatching(entry.prepared.source).catch((error) => {
							debugLog("Image matching failed; retaining the source attachment", error);
							return undefined;
						})
						: Promise.resolve(undefined),
				),
			);
			for (let candidateIndex = 0; candidateIndex < normalizedCandidates.length; candidateIndex += 1) {
				const normalized = normalizedCandidates[candidateIndex];
				if (!normalized) continue;
				const key = `${normalized.mimeType}\u0000${normalized.data}`;
				existingIndexes[candidateIndex] = existingByContent.get(key)?.shift();
			}
		}
		const preparedImages = await Promise.all(
			candidates.map(async ({ entry }, candidateIndex) => ({
				image: await entry.prepared.model(),
				existingIndex: existingIndexes[candidateIndex],
			})),
		);
		const references = await Promise.all(
			candidates.map(({ entry }) => entry.prepared.persist()),
		);
		let transformedText = fullText;
		for (const [index, candidate] of candidates.entries()) {
			const reference = references[index];
			const marker = reference
				? createImageMarkerLink(candidate.entry.placeholder, reference)
				: candidate.entry.placeholder;
			transformedText = transformedText.replace(candidate.token, marker);
		}
		const images = [...(event.images ?? [])];
		for (const { image, existingIndex } of preparedImages) {
			if (existingIndex === undefined) images.push(image);
			else images[existingIndex] = image;
		}
		if (inputSession === sessionGeneration) {
			for (const { entry } of candidates) {
				if (draft.has(entry)) draft.images.delete(entry.placeholder);
			}
			scanGeneration += 1;
			lastScannedText = undefined;
			refreshWidget(ctx);
			editorAdapter.refresh();
		}
		return { action: "transform", text: transformedText, images };
	});
}