/**
 * see — Vision proxy extension for pi
 *
 * Allows ANY model in pi to "see" — routes image+prompt requests through
 * a configurable vision-capable model and returns the result as text.
 *
 * Smart routing:
 *   1. If user set a vision model via /seemodel → proxies through that
 *   2. Else if current model has native vision → uses it directly
 *   3. Else → prompts user to configure a proxy model
 *
 * The LLM gets a `see` tool it can call autonomously. Users also have the
 * `/see` command for manual use. Can find and read images from disk.
 *
 * Usage:
 *   /see describe the latest screenshot
 *   /see what's in ~/path/to/image.png
 *   /see model
 *
 * Install: copy to ~/.pi/agent/extensions/ or .pi/extensions/ then /reload
 */

import { readFileSync } from "node:fs";
import { readdir, stat } from "node:fs/promises";
import { homedir } from "node:os";
import { extname, join, resolve } from "node:path";
import { complete } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

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

const IMAGE_TYPES = new Map([
	[".png", "image/png"],
	[".jpg", "image/jpeg"],
	[".jpeg", "image/jpeg"],
	[".gif", "image/gif"],
	[".webp", "image/webp"],
	[".bmp", "image/bmp"],
]);

const DEFAULT_SCREENSHOT_DIR = join(homedir(), "Pictures", "Screenshots");
const FALLBACK_SCREENSHOT_DIR = join(homedir(), "Desktop");

// ── State persisted to session ──────────────────────────────────────────────

interface SeeState {
	defaultModel?: string;
}

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

function expandPath(p: string): string {
	if (p.startsWith("~/")) return join(homedir(), p.slice(2));
	if (p === "~") return homedir();
	return resolve(p);
}

function isImageFile(p: string): boolean {
	return IMAGE_TYPES.has(extname(p).toLowerCase());
}

function readImageBase64(filePath: string): string {
	return readFileSync(filePath).toString("base64");
}

function mimeForFile(filePath: string): string {
	return IMAGE_TYPES.get(extname(filePath).toLowerCase()) ?? "image/png";
}

/**
 * Find the most recently modified image in a directory.
 */
async function findLatestImage(dir: string): Promise<{ path: string; mimeType: string } | null> {
	const results = await findLatestImages(dir, 1);
	return results.length > 0 ? results[0] : null;
}

/**
 * Find the N most recently modified images in a directory.
 */
async function findLatestImages(dir: string, count: number = 1): Promise<Array<{ path: string; mimeType: string }>> {
	let files: string[];
	try {
		files = await readdir(dir);
	} catch {
		return [];
	}

	const imageFiles = files.filter((f) => isImageFile(f));
	if (imageFiles.length === 0) return [];

	const withStats = await Promise.all(
		imageFiles.map(async (f) => {
			const fullPath = join(dir, f);
			try {
				const s = await stat(fullPath);
				return { path: fullPath, mtime: s.mtimeMs, ext: extname(f).toLowerCase() };
			} catch {
				return null;
			}
		}),
	);

	const valid = withStats.filter(Boolean) as Array<{ path: string; mtime: number; ext: string }>;
	valid.sort((a, b) => b.mtime - a.mtime);
	return valid.slice(0, count).map((f) => ({ path: f.path, mimeType: mimeForFile(f.path) }));
}

/**
 * Parse a natural-language instruction to discover image files.
 */
async function resolveImages(
	instruction: string,
): Promise<{
	prompt: string;
	images: Array<{ data: string; mimeType: string }>;
	sourceLabel?: string;
} | null> {
	// Case 1: instruction ends with an explicit file path
	const fileEndMatch = instruction.match(/^(.*?)\s+((?:\/[^\s]+|~[^\s]+|\.[\/\\][^\s]+))$/);
	if (fileEndMatch) {
		const prompt = fileEndMatch[1].trim();
		const rawPath = fileEndMatch[2].trim();
		const expanded = expandPath(rawPath);

		try {
			const s = await stat(expanded);
			if (s.isFile() && isImageFile(expanded)) {
				const data = readImageBase64(expanded);
				return { prompt: prompt || "Describe this image.", images: [{ data, mimeType: mimeForFile(expanded) }], sourceLabel: expanded };
			}
			if (s.isDirectory()) {
				const found = await findLatestImage(expanded);
				if (found) {
					const data = readImageBase64(found.path);
					return { prompt: prompt || "Describe this image.", images: [{ data, mimeType: found.mimeType }], sourceLabel: found.path };
				}
				return { prompt: `${prompt}\n\n(I tried to find images in ${expanded} but none were found)`, images: [], sourceLabel: undefined };
			}
		} catch {
			// Not a valid path — fall through
		}
	}

	// Case 2: instruction references a directory with screenshot/latest
	const dirMatch = instruction.match(
		/(?:latest|newest|most\s+recent)?\s*(?:screenshot|image|picture|photo)?\s*(?:in|from|at|of)?\s*((?:\/[^\s]+|~[^\s]+|\.[\/\\][^\s]+))$/i,
	);
	if (dirMatch) {
		const rawDir = dirMatch[1].trim();
		const expanded = expandPath(rawDir);
		const prompt = instruction.replace(dirMatch[0], "").trim();

		const found = await findLatestImage(expanded);
		if (found) {
			const data = readImageBase64(found.path);
			return { prompt: prompt || "Describe this image.", images: [{ data, mimeType: found.mimeType }], sourceLabel: found.path };
		}

		try {
			await stat(expanded);
			return { prompt: `${prompt}\n\n(I looked in ${expanded} but found no image files)`, images: [], sourceLabel: undefined };
		} catch {
			return { prompt: `${prompt}\n\n(I tried to look in ${expanded} but the directory doesn't exist)`, images: [], sourceLabel: undefined };
		}
	}

	// Case 3: "screenshot" mentioned without a directory — use default
	if (/\b(screenshot|snapshot)\b/i.test(instruction)) {
		// Parse optional count: "latest 3 screenshots", "latest three screenshots", "last 2"
		const numberWords: Record<string, number> = {
			one: 1, two: 2, three: 3, four: 4, five: 5,
			six: 6, seven: 7, eight: 8, nine: 9, ten: 10,
		};
		let count = 1;
		const countMatch = instruction.match(/\b(latest|last|newest)\s+(\d+|\w+)\s+/i);
		if (countMatch) {
			const numStr = countMatch[2].toLowerCase();
			count = parseInt(numStr, 10);
			if (isNaN(count)) count = numberWords[numStr] ?? 1;
			if (count < 1) count = 1;
			if (count > 20) count = 20;
		}

		const prompt = instruction.trim() || "Describe these images.";

		let results = await findLatestImages(DEFAULT_SCREENSHOT_DIR, count);
		if (results.length === 0) {
			results = await findLatestImages(FALLBACK_SCREENSHOT_DIR, count);
		}

		if (results.length > 0) {
			const images = results.map((r) => ({ data: readImageBase64(r.path), mimeType: r.mimeType }));
			const label = results.length === 1 ? results[0].path : results.map((r) => r.path).join(", ");
			return { prompt, images, sourceLabel: label };
		}

		return {
			prompt: `${prompt}\n\n(I looked in ${DEFAULT_SCREENSHOT_DIR} and ${FALLBACK_SCREENSHOT_DIR} but found no image files)`,
			images: [],
			sourceLabel: undefined,
		};
	}

	// Case 4: just a file path without leading prompt
	const barePathMatch = instruction.match(/^((?:\/[^\s]+|~[^\s]+|\.[\/\\][^\s]+))$/);
	if (barePathMatch) {
		const rawPath = barePathMatch[1].trim();
		const expanded = expandPath(rawPath);
		try {
			const s = await stat(expanded);
			if (s.isFile() && isImageFile(expanded)) {
				const data = readImageBase64(expanded);
				return { prompt: "Describe this image.", images: [{ data, mimeType: mimeForFile(expanded) }], sourceLabel: expanded };
			}
		} catch {
			// Not valid
		}
	}

	return null;
}

/**
 * Resolve which model to use for vision.
 *
 * Priority:
 * 1. If user explicitly configured a vision model via /seemodel → use that
 * 2. Else if the current conversation model has native vision → use it directly
 * 3. Else → error, tell user to set up /seemodel
 */
async function resolveVisionModel(
	ctx: ExtensionContext,
	explicitModel: string | null,
): Promise<{
	model: NonNullable<ReturnType<ExtensionContext["modelRegistry"]["find"]>>;
	auth: { apiKey: string; headers?: Record<string, string> };
	isNative: boolean;
} | { error: string }> {
	// Priority 1: user explicitly configured a vision model
	if (explicitModel) {
		const sepIndex = explicitModel.indexOf("/");
		if (sepIndex !== -1) {
			const provider = explicitModel.slice(0, sepIndex);
			const modelId = explicitModel.slice(sepIndex + 1);
			const model = ctx.modelRegistry.find(provider, modelId);
			if (model) {
				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
				if (auth.ok && auth.apiKey) {
					return { model, auth: { apiKey: auth.apiKey, headers: auth.headers }, isNative: false };
				}
			}
		}
	}

	// Priority 2: current conversation model has native vision
	const currentModel = ctx.model;
	if (currentModel?.input?.includes("image")) {
		const auth = await ctx.modelRegistry.getApiKeyAndHeaders(currentModel);
		if (auth.ok && auth.apiKey) {
			return { model: currentModel, auth: { apiKey: auth.apiKey, headers: auth.headers }, isNative: true };
		}
	}

	// Priority 3: nothing works
	const hint = explicitModel
		? `Model "${explicitModel}" not found or has no API key.`
		: "Your current model does not support vision. Use `/seemodel` to set a vision-capable proxy model.";
	return { error: hint };
}

/**
 * Call the vision model with a prompt and images.
 * Returns the text result.
 */
async function callVisionModel(
	prompt: string,
	images: Array<{ data: string; mimeType: string }> | undefined,
	model: NonNullable<ReturnType<ExtensionContext["modelRegistry"]["find"]>>,
	auth: { apiKey: string; headers?: Record<string, string> },
	signal?: AbortSignal,
): Promise<{ ok: true; result: string } | { ok: false; error: string }> {
	try {
		const content: Array<{ type: "text" | "image"; text?: string; data?: string; mimeType?: string }> = [
			{ type: "text", text: prompt },
		];

		if (images) {
			for (const img of images) {
				content.push({ type: "image", data: img.data, mimeType: img.mimeType });
			}
		}

		const response = await complete(
			model,
			{
				messages: [{ role: "user", content, timestamp: Date.now() }],
			},
			{ apiKey: auth.apiKey, headers: auth.headers, signal },
		);

		const result = response.content
			.filter((c): c is { type: "text"; text: string } => c.type === "text")
			.map((c) => c.text)
			.join("\n");

		return { ok: true, result };
	} catch (err: unknown) {
		const message = err instanceof Error ? err.message : String(err);
		return { ok: false, error: message };
	}
}

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

export default function seeExtension(pi: ExtensionAPI) {
	let defaultModel: string | null = null;

	// ── Status display ────────────────────────────────────────────────────

	function updateSeeStatus(ctx: ExtensionContext) {
		const current = ctx.model;
		if (defaultModel) {
			ctx.ui.setStatus("see", `👁 ${defaultModel}`);
		} else if (current?.input?.includes("image")) {
			ctx.ui.setStatus("see", `👁 ${current.provider}/${current.id} (native)`);
		} else {
			ctx.ui.setStatus("see", undefined);
		}
	}

	// ── State persistence ──────────────────────────────────────────────────

	function restoreState(ctx: ExtensionContext) {
		const entries = ctx.sessionManager.getBranch();
		for (const entry of entries) {
			if (entry.type === "custom" && entry.customType === "see-config") {
				const data = entry.data as SeeState | undefined;
				if (data?.defaultModel) {
					defaultModel = data.defaultModel;
				}
			}
		}
		updateSeeStatus(ctx);
	}

	// Ensure the `see` tool is always in the active tool set
	function activateSeeTool() {
		const active = pi.getActiveTools();
		if (!active.includes("see")) {
			pi.setActiveTools([...active, "see"]);
		}
	}

	pi.on("session_start", async (_event, ctx) => {
		restoreState(ctx);
		activateSeeTool();
	});

	pi.on("session_tree", async (_event, ctx) => {
		restoreState(ctx);
		activateSeeTool();
	});

	// ── Model selector ────────────────────────────────────────────────────

	async function showModelSelector(ctx: ExtensionContext) {
		const available = ctx.modelRegistry.getAvailable();
		const visionModels = available.filter((m) => m.input?.includes("image"));

		if (visionModels.length === 0) {
			ctx.ui.notify("No vision-capable models with configured auth found", "warning");
			return;
		}

		const items = visionModels.map((m) => `${m.provider}/${m.id}`);
		const selected = await ctx.ui.select("Select default vision model", items);

		if (selected) {
			defaultModel = selected;
			updateSeeStatus(ctx);
			pi.appendEntry<SeeState>("see-config", { defaultModel });
			ctx.ui.notify(`Default vision model set to: ${defaultModel}`, "success");
		}
	}

	pi.registerCommand("seemodel", {
		description: "Select the default vision model for /see",
		handler: async (_args, ctx) => {
			await showModelSelector(ctx);
		},
	});

	// ── User-facing /see processing ─────────────────────────────────────────

	async function handleSeeForUser(
		prompt: string,
		images: Array<{ data: string; mimeType: string }> | undefined,
		ctx: ExtensionContext,
		sourceLabel?: string,
	) {
		const resolved = await resolveVisionModel(ctx, defaultModel);
		if ("error" in resolved) {
			ctx.ui.notify(resolved.error, "warning");
			return;
		}
		const { model, auth, isNative } = resolved;

		const nativeTag = isNative ? " (native)" : "";
		const label = sourceLabel ? ` from ${sourceLabel}` : "";
		ctx.ui.notify(`👁  ${model.provider}/${model.id}${nativeTag}${label}...`, "info");

		const visionResult = await callVisionModel(prompt, images, model, auth, ctx.signal);

		if (!visionResult.ok) {
			ctx.ui.notify(`Error: ${visionResult.error}`, "error");
			return;
		}

		// Show as notification
		ctx.ui.notify(visionResult.result, "info");

		// Inject into conversation as a permanent text-only message
		const sourceInfo = sourceLabel
			? `**Source:** ${sourceLabel}\n`
			: images && images.length > 0
				? `**Source:** attached image(s)\n`
				: "";

		const nativeTag2 = isNative ? " — native vision" : "";
		const summary = [
			`**👁 See** (${model.provider}/${model.id}${nativeTag2})`,
			sourceInfo,
			`**Prompt:** ${prompt}`,
			"",
			visionResult.result,
		].filter(Boolean).join("\n");

		pi.sendMessage({
			customType: "see-result",
			content: [{ type: "text", text: summary }],
			display: true,
		});
	}

	// ── Tool: see (callable by the LLM) ─────────────────────────────────────

	pi.registerTool({
		name: "see",
		label: "See (vision proxy)",
		description:
			"Analyze any image, screenshot, photo, diagram, chart, drawing, UI mockup, or scanned " +
			"document. Extracts text, describes content, interprets visual information, identifies " +
			"objects, reads UI elements, translates text in images, analyzes layouts, and more. " +
			"Provide an imagePath or let it find the latest screenshot automatically. " +
			"If the current model already has native vision, it uses that directly. " +
			"Use this tool any time you need to understand visual content — never guess or say you " +
			"cannot process images without trying this tool first.",
		promptGuidelines: [
			"You have a `see` tool that can analyze images, screenshots, photos, diagrams, charts, " +
			"drawings, UI mockups, scanned documents, or any visual content. " +
			"If your current model already supports vision, the `see` tool uses your own vision " +
			"capabilities directly. Otherwise it proxies through a configured vision-capable model. " +
			"Use the `see` tool whenever the conversation involves: " +
			"attached images or screenshots, questions about visual content, requests to read text " +
			"from images, describe photos, analyze UI layouts, interpret charts or diagrams, " +
			"extract information from scanned documents, review designs or mockups, identify objects " +
			"in photos, or any task that requires understanding visual information. " +
			"Always prefer using the `see` tool over guessing or saying you cannot process images. " +
			"Be specific in your prompt about what you need to know — don't just say 'describe this', " +
			"tell it exactly what you're looking for.",
			"You can pass multiple image paths via imagePaths[] to compare images side by side or " +
			"analyze them together.",
		],
		parameters: Type.Object({
			prompt: Type.String({
				description:
					"What to look for or ask about in the image. Be specific about what you need to know — " +
					"e.g. 'describe this scene', 'read all the text', 'what color is the button?', " +
					"'what are the dimensions of the chart?', 'translate the text in this image', " +
					"'identify the objects on the desk', 'explain this diagram', 'what error message is shown?', " +
					"'list the menu items visible', 'describe the UI layout'. The more specific your " +
					"prompt, the better the result.",
			}),
			imagePaths: Type.Optional(Type.Array(Type.String(), {
				description:
					"One or more paths to image files (e.g., ~/Pictures/screenshot.png). " +
					"Pass multiple to compare or analyze together. " +
					"If omitted, automatically finds the latest screenshot in ~/Pictures/Screenshots/.",
			})),
		}),
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			const resolved = await resolveVisionModel(ctx, defaultModel);
			if ("error" in resolved) {
				return {
					content: [{ type: "text", text: resolved.error }],
					isError: true,
				};
			}
			const { model, auth, isNative } = resolved;

			// Resolve images from paths or find latest screenshot
			let images: Array<{ data: string; mimeType: string }> = [];
			let sourceLabels: string[] = [];

			if (params.imagePaths && params.imagePaths.length > 0) {
				for (const rawPath of params.imagePaths) {
					const expanded = expandPath(rawPath);
					try {
						const s = await stat(expanded);
						if (s.isFile() && isImageFile(expanded)) {
							images.push({ data: readImageBase64(expanded), mimeType: mimeForFile(expanded) });
							sourceLabels.push(expanded);
						} else if (s.isDirectory()) {
							const found = await findLatestImage(expanded);
							if (found) {
								images.push({ data: readImageBase64(found.path), mimeType: found.mimeType });
								sourceLabels.push(found.path);
							}
						}
					} catch {
						// Path doesn't exist — skip
					}
				}
			}

			if (images.length === 0) {
				// No paths or none resolved — try latest screenshot
				const found = await findLatestImage(DEFAULT_SCREENSHOT_DIR) ?? await findLatestImage(FALLBACK_SCREENSHOT_DIR);
				if (found) {
					images.push({ data: readImageBase64(found.path), mimeType: found.mimeType });
					sourceLabels.push(found.path);
				}
			}

			if (images.length === 0) {
				return {
					content: [{ type: "text", text: "Could not find any images to analyze. Provide valid imagePaths or ensure screenshots exist." }],
					isError: true,
				};
			}

			const visionResult = await callVisionModel(params.prompt, images, model, auth, signal);

			if (!visionResult.ok) {
				return {
					content: [{ type: "text", text: `Vision model error: ${visionResult.error}` }],
					isError: true,
				};
			}

			const labelStr = sourceLabels.length > 0
				? `[Analyzed: ${sourceLabels.join(", ")}]\n\n`
				: "";
			const result = `${labelStr}${visionResult.result}`;

			return {
				content: [{ type: "text", text: result }],
			};
		},
	});

	// ── Input handler (catches /see for the user) ───────────────────────────

	pi.on("input", async (event, ctx) => {
		const text = event.text.trim();
		if (!text.startsWith("/see")) return { action: "continue" };

		if (text === "/see model") {
			await showModelSelector(ctx);
			return { action: "handled" };
		}

		const instruction = text.replace(/^\/see\s*/, "").trim();
		const hasAttachedImages = event.images && event.images.length > 0;

		if (!instruction && !hasAttachedImages) {
			ctx.ui.notify(
				"Usage: /see <prompt> (attach images, provide image paths, or describe what to find)",
				"warning",
			);
			return { action: "handled" };
		}

		const resolved = instruction ? await resolveImages(instruction) : null;

		if (resolved) {
			await handleSeeForUser(resolved.prompt, resolved.images, ctx, resolved.sourceLabel);
		} else if (hasAttachedImages) {
			await handleSeeForUser(
				instruction || "Describe these images.",
				event.images?.map((img) => ({ data: img.data, mimeType: img.mimeType })),
				ctx,
			);
		} else {
			await handleSeeForUser(instruction, undefined, ctx);
		}

		return { action: "handled" };
	});
}