/**
 * copy-code: copy fenced code blocks from recent assistant messages to the
 * system clipboard.
 *
 * Triggers:
 *   - shortcut: ctrl+shift+y
 *   - command:  /cc            -> picker over recent blocks
 *   - command:  /cc <n>        -> copy the n-th most-recent block (1-based)
 *   - command:  /cc last       -> copy the most recent block
 *   - command:  /cc all        -> concatenate all blocks from the latest
 *                                 assistant message
 *
 * Sources scanned (newest first): assistant text content and toolResult text
 * content. We deliberately ignore thinking content.
 */

import { spawn } from "node:child_process";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
	SessionEntry,
	SessionMessageEntry,
} from "@mariozechner/pi-coding-agent";
import type { AssistantMessage, TextContent, ToolResultMessage } from "@mariozechner/pi-ai";

interface CodeBlock {
	lang: string;
	code: string;
	lines: number;
	source: "assistant" | "toolResult";
	toolName?: string;
	timestamp: number;
}

interface ClipboardCmd {
	cmd: string;
	args: string[];
}

const FENCE_RE = /(^|\n)([ \t]{0,3})(`{3,}|~{3,})([^\n`]*)\n([\s\S]*?)\n\2\3[ \t]*(?=\n|$)/g;

function extractFromText(text: string): Array<{ lang: string; code: string }> {
	const blocks: Array<{ lang: string; code: string }> = [];
	FENCE_RE.lastIndex = 0;
	let m: RegExpExecArray | null;
	// biome-ignore lint/suspicious/noAssignInExpressions: standard regex iteration
	while ((m = FENCE_RE.exec(text)) !== null) {
		const lang = (m[4] ?? "").trim().split(/\s+/)[0] ?? "";
		const code = m[5] ?? "";
		if (code.length > 0) blocks.push({ lang, code });
	}
	return blocks;
}

function textOfContent(content: string | (TextContent | { type: string })[]): string {
	if (typeof content === "string") return content;
	return content
		.filter((c): c is TextContent => c.type === "text")
		.map((c) => c.text)
		.join("\n");
}

function collectBlocks(entries: readonly SessionEntry[]): CodeBlock[] {
	const out: CodeBlock[] = [];
	// Walk newest -> oldest so the picker shows recent first.
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry.type !== "message") continue;
		const msg = (entry as SessionMessageEntry).message;
		const ts = typeof entry.timestamp === "string" ? Date.parse(entry.timestamp) : Date.now();

		if (msg.role === "assistant") {
			const a = msg as AssistantMessage;
			const text = a.content
				.filter((c): c is TextContent => c.type === "text")
				.map((c) => c.text)
				.join("\n\n");
			for (const b of extractFromText(text)) {
				out.push({
					lang: b.lang,
					code: b.code,
					lines: b.code.split("\n").length,
					source: "assistant",
					timestamp: ts,
				});
			}
		} else if (msg.role === "toolResult") {
			const t = msg as ToolResultMessage;
			const text = textOfContent(t.content);
			for (const b of extractFromText(text)) {
				out.push({
					lang: b.lang,
					code: b.code,
					lines: b.code.split("\n").length,
					source: "toolResult",
					toolName: t.toolName,
					timestamp: ts,
				});
			}
		}
	}
	return out;
}

function blocksFromLatestAssistant(entries: readonly SessionEntry[]): CodeBlock[] {
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry.type !== "message") continue;
		const msg = (entry as SessionMessageEntry).message;
		if (msg.role !== "assistant") continue;
		const text = (msg as AssistantMessage).content
			.filter((c): c is TextContent => c.type === "text")
			.map((c) => c.text)
			.join("\n\n");
		const ts = typeof entry.timestamp === "string" ? Date.parse(entry.timestamp) : Date.now();
		return extractFromText(text).map((b) => ({
			lang: b.lang,
			code: b.code,
			lines: b.code.split("\n").length,
			source: "assistant" as const,
			timestamp: ts,
		}));
	}
	return [];
}

async function detectClipboard(): Promise<ClipboardCmd | null> {
	const candidates: ClipboardCmd[] = [];
	if (process.platform === "darwin") {
		candidates.push({ cmd: "pbcopy", args: [] });
	} else if (process.platform === "win32") {
		candidates.push({ cmd: "clip.exe", args: [] });
	} else {
		// Linux / WSL. Prefer Wayland when available.
		if (process.env.WAYLAND_DISPLAY) {
			candidates.push({ cmd: "wl-copy", args: [] });
		}
		candidates.push({ cmd: "xclip", args: ["-selection", "clipboard"] });
		candidates.push({ cmd: "xsel", args: ["--clipboard", "--input"] });
		// WSL fallback
		candidates.push({ cmd: "clip.exe", args: [] });
	}

	for (const c of candidates) {
		if (await hasCommand(c.cmd)) return c;
	}
	return null;
}

function hasCommand(cmd: string): Promise<boolean> {
	return new Promise((resolve) => {
		const p = spawn("sh", ["-c", `command -v ${JSON.stringify(cmd)}`], { stdio: "ignore" });
		p.on("exit", (code) => resolve(code === 0));
		p.on("error", () => resolve(false));
	});
}

function writeClipboard(clip: ClipboardCmd, data: string): Promise<void> {
	return new Promise((resolve, reject) => {
		const p = spawn(clip.cmd, clip.args, { stdio: ["pipe", "ignore", "pipe"] });
		let stderr = "";
		p.stderr.on("data", (chunk) => {
			stderr += chunk.toString();
		});
		p.on("error", reject);
		p.on("exit", (code) => {
			if (code === 0) resolve();
			else reject(new Error(`${clip.cmd} exited ${code}: ${stderr.trim()}`));
		});
		p.stdin.end(data);
	});
}

function preview(code: string, max = 60): string {
	const firstLine = code.split("\n").find((l) => l.trim().length > 0) ?? "";
	const stripped = firstLine.trim();
	return stripped.length > max ? `${stripped.slice(0, max - 1)}...` : stripped;
}

function blockLabel(b: CodeBlock, idx: number): string {
	const where = b.source === "toolResult" ? `${b.toolName ?? "tool"}` : "assistant";
	const lang = b.lang || "text";
	return `${idx + 1}. [${where}] ${lang} (${b.lines}L) ${preview(b.code)}`;
}

export default function (pi: ExtensionAPI) {
	let clipboard: ClipboardCmd | null = null;
	let detected = false;

	const ensureClipboard = async (ctx: ExtensionContext): Promise<ClipboardCmd | null> => {
		if (!detected) {
			clipboard = await detectClipboard();
			detected = true;
		}
		if (!clipboard) {
			ctx.ui.notify(
				"copy-code: no clipboard tool found (install wl-copy, xclip, or xsel)",
				"error",
			);
		}
		return clipboard;
	};

	const copy = async (ctx: ExtensionContext, code: string, label: string) => {
		const clip = await ensureClipboard(ctx);
		if (!clip) return;
		try {
			await writeClipboard(clip, code);
			ctx.ui.notify(`Copied ${label} (${code.length} chars) via ${clip.cmd}`, "success");
		} catch (err) {
			ctx.ui.notify(
				`copy-code: clipboard write failed: ${err instanceof Error ? err.message : String(err)}`,
				"error",
			);
		}
	};

	const pickAndCopy = async (ctx: ExtensionContext) => {
		const blocks = collectBlocks(ctx.sessionManager.getEntries()).slice(0, 30);
		if (blocks.length === 0) {
			ctx.ui.notify("copy-code: no code blocks found in this session", "warning");
			return;
		}
		if (blocks.length === 1) {
			await copy(ctx, blocks[0].code, `${blocks[0].lang || "text"} block`);
			return;
		}
		const choice = await ctx.ui.select(
			"Copy which code block?",
			blocks.map((b, i) => blockLabel(b, i)),
		);
		if (choice == null) return;
		// ctx.ui.select returns the selected string; map back by index prefix.
		const idx = blocks.findIndex((b, i) => blockLabel(b, i) === choice);
		const block = idx >= 0 ? blocks[idx] : blocks[0];
		await copy(ctx, block.code, `${block.lang || "text"} block`);
	};

	pi.on("session_start", async () => {
		clipboard = await detectClipboard();
		detected = true;
	});

	pi.registerShortcut("ctrl+shift+y", {
		description: "Copy a code block from recent messages",
		handler: async (ctx) => {
			await pickAndCopy(ctx);
		},
	});

	pi.registerCommand("cc", {
		description: "Copy a code block. Usage: /cc [n|last|all]",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			const arg = args.trim().toLowerCase();

			if (arg === "" || arg === "pick") {
				await pickAndCopy(ctx);
				return;
			}

			if (arg === "all") {
				const blocks = blocksFromLatestAssistant(ctx.sessionManager.getEntries());
				if (blocks.length === 0) {
					ctx.ui.notify("copy-code: no code blocks in latest assistant message", "warning");
					return;
				}
				const joined = blocks.map((b) => b.code).join("\n\n");
				await copy(ctx, joined, `${blocks.length} blocks`);
				return;
			}

			const blocks = collectBlocks(ctx.sessionManager.getEntries());
			if (blocks.length === 0) {
				ctx.ui.notify("copy-code: no code blocks found", "warning");
				return;
			}

			let idx: number;
			if (arg === "last") {
				idx = 0;
			} else {
				const n = Number.parseInt(arg, 10);
				if (!Number.isFinite(n) || n < 1) {
					ctx.ui.notify(`copy-code: invalid argument '${args}'`, "error");
					return;
				}
				idx = n - 1;
			}

			if (idx >= blocks.length) {
				ctx.ui.notify(
					`copy-code: only ${blocks.length} block(s) available`,
					"warning",
				);
				return;
			}
			const block = blocks[idx];
			await copy(ctx, block.code, `${block.lang || "text"} block (#${idx + 1})`);
		},
	});
}