import type {
	ExtensionAPI,
	ExtensionContext,
	Theme,
} from "@earendil-works/pi-coding-agent";
import {
	Input,
	Key,
	matchesKey,
	truncateToWidth,
	visibleWidth,
	type Component,
	type Focusable,
	type KeyId,
	type TUI,
} from "@earendil-works/pi-tui";

interface Annotation {
	start: number; // source line index, inclusive
	end: number; // source line index, inclusive
	note: string;
}

interface Draft {
	annotations: Annotation[];
	cursor: number;
}

type Mode = "nav" | "span" | "note";

interface Row {
	line: number; // source line this wrapped segment belongs to
	text: string;
}

const GUTTER = 4;

// Quotes exist to locate a passage in the agent's own message, not to
// duplicate it - spans past these bounds are anchored by their edges.
const MAX_QUOTE_LINES = 4;
const MAX_QUOTE_CHARS = 300;

/**
 * Word wrap that hard-splits words longer than width. Rows built from this are
 * display-only: quotes sent to the agent are taken from the source lines, so
 * wrapping can never corrupt the anchor text.
 */
function wrapLine(line: string, width: number): string[] {
	const expanded = line.replace(/\t/g, "   ");
	if (expanded === "") return [""];
	if (width <= 0) return [expanded];

	const out: string[] = [];
	let cur = "";
	for (const word of expanded.split(" ")) {
		if (visibleWidth(word) > width) {
			if (cur) {
				out.push(cur);
				cur = "";
			}
			let rest = word;
			while (visibleWidth(rest) > width) {
				const part = truncateToWidth(rest, width, "");
				out.push(part);
				rest = rest.slice(part.length);
			}
			cur = rest;
			continue;
		}
		const candidate = cur ? `${cur} ${word}` : word;
		if (visibleWidth(candidate) <= width) {
			cur = candidate;
		} else {
			if (cur) out.push(cur);
			cur = word;
		}
	}
	if (cur || out.length === 0) out.push(cur);
	return out;
}

export function formatAnnotations(
	annotations: Annotation[],
	messageText: string,
): string {
	const lines = messageText.replace(/\r\n/g, "\n").split("\n");
	let truncated = false;
	const body = annotations
		.map((a, i) => {
			const span = lines.slice(a.start, a.end + 1);
			const totalChars = span.reduce((n, l) => n + l.length, 0);
			let quoteLines: string[];
			if (span.length <= MAX_QUOTE_LINES && totalChars <= MAX_QUOTE_CHARS) {
				quoteLines = span;
			} else if (span.length >= 3) {
				truncated = true;
				quoteLines = [...span.slice(0, 2), "…", span[span.length - 1]!];
			} else {
				truncated = true;
				quoteLines = span.map((l) =>
					l.length > 200 ? `${l.slice(0, 200)} …` : l,
				);
			}
			const quote = quoteLines.map((l) => `> ${l}`).join("\n");
			return `[${i + 1}]\n${quote}\nNote: ${a.note}`;
		})
		.join("\n\n");
	const header =
		`${annotations.length} comment(s) on your last message - each quote anchors the comment below it. ` +
		"Address each by number and flag any conflicts." +
		(truncated
			? " Long quotes are truncated with …; the full text is in your message above."
			: "");
	return `${header}\n\n${body}`;
}

export class AnnotateOverlay implements Component, Focusable {
	private _focused = false;
	get focused() {
		return this._focused;
	}
	set focused(value: boolean) {
		this._focused = value;
		this.noteInput.focused = value;
	}

	private tui: TUI;
	private theme: Theme;
	private lines: string[];
	private done: (result: Annotation[] | null) => void;

	private mode: Mode = "nav";
	private cursor = 0;
	private anchor: number | null = null;
	private editing: number | null = null;
	private annotations: Annotation[] = [];
	private noteInput = new Input();
	private hint: string | null = null;

	private rows: Row[] = [];
	private cachedWidth?: number;
	private scrollOffset = 0;
	private lastBodyHeight = 10;

	constructor(
		tui: TUI,
		theme: Theme,
		messageText: string,
		done: (result: Annotation[] | null) => void,
		draft?: Draft,
	) {
		this.tui = tui;
		this.theme = theme;
		this.lines = messageText.replace(/\r\n/g, "\n").split("\n");
		this.done = done;
		if (draft) {
			this.annotations = draft.annotations.map((a) => ({ ...a }));
			this.cursor = Math.min(
				Math.max(0, draft.cursor),
				this.lines.length - 1,
			);
		}
		this.noteInput.onSubmit = (value) => this.commitNote(value);
		this.noteInput.onEscape = () => this.cancelNote();
	}

	getDraftState(): Draft {
		return {
			annotations: this.annotations.map((a) => ({ ...a })),
			cursor: this.cursor,
		};
	}

	handleInput(data: string): void {
		if (this.mode === "note") {
			this.noteInput.handleInput(data);
		} else if (matchesKey(data, Key.escape)) {
			this.handleEscape();
		} else if (matchesKey(data, Key.up) || data === "k") {
			this.moveCursor(-1);
		} else if (matchesKey(data, Key.down) || data === "j") {
			this.moveCursor(1);
		} else if (matchesKey(data, Key.pageUp)) {
			this.moveCursor(-this.lastBodyHeight);
		} else if (matchesKey(data, Key.pageDown)) {
			this.moveCursor(this.lastBodyHeight);
		} else if (data === "g") {
			this.cursor = 0;
		} else if (data === "G") {
			this.cursor = this.lines.length - 1;
		} else if (this.mode === "nav" && data === "v") {
			this.anchor = this.cursor;
			this.mode = "span";
			this.hint = null;
		} else if (
			this.mode === "nav" &&
			(data === "S" || matchesKey(data, "shift+enter"))
		) {
			// S is the send fallback for terminals without the Kitty keyboard
			// protocol, where shift+enter arrives indistinguishable from enter.
			if (this.annotations.length > 0) {
				this.done(this.annotations);
				return;
			}
			this.hint = "no annotations yet - mark a span with v or press i";
		} else if (
			this.mode === "nav" &&
			(data === "i" || data === "a" || matchesKey(data, Key.enter))
		) {
			const annIdx = this.annotationAt(this.cursor);
			if (annIdx >= 0) {
				this.openEdit(annIdx);
			} else {
				this.anchor = this.cursor;
				this.openNote();
			}
		} else if (
			this.mode === "span" &&
			(data === "s" || data === "v" || matchesKey(data, Key.enter))
		) {
			this.openNote();
		} else if (this.mode === "nav" && (data === "x" || data === "d")) {
			const annIdx = this.annotationAt(this.cursor);
			if (annIdx >= 0) {
				this.annotations.splice(annIdx, 1);
			} else {
				this.hint = "no annotation on this line";
			}
		} else if (this.mode === "nav" && data === "C") {
			if (this.annotations.length > 0) {
				this.hint = `cleared ${this.annotations.length} annotation(s)`;
				this.annotations = [];
			}
		}
		this.tui.requestRender();
	}

	render(width: number): string[] {
		const theme = this.theme;
		const rows = this.buildRows(width);

		const chrome = 6 + (this.mode === "note" ? 2 : 0);
		const maxBody = Math.max(
			3,
			Math.floor(this.tui.terminal.rows * 0.8) - chrome,
		);
		const bodyHeight = Math.max(3, Math.min(maxBody, rows.length));
		this.lastBodyHeight = bodyHeight;

		const cursorRow = rows.findIndex((r) => r.line === this.cursor);
		if (cursorRow < this.scrollOffset) this.scrollOffset = cursorRow;
		if (cursorRow >= this.scrollOffset + bodyHeight)
			this.scrollOffset = cursorRow - bodyHeight + 1;
		this.scrollOffset = Math.max(
			0,
			Math.min(this.scrollOffset, rows.length - bodyHeight),
		);

		const span = this.spanRange();
		const out: string[] = [];

		const inner = width - 2;
		const borderV = theme.fg("accent", theme.bold("┃"));
		const padInner = (content: string): string => {
			const t = truncateToWidth(content, inner);
			return t + " ".repeat(Math.max(0, inner - visibleWidth(t)));
		};
		const boxLine = (content: string): string =>
			borderV + padInner(content) + borderV;

		out.push(theme.fg("accent", theme.bold(`┏${"━".repeat(inner)}┓`)));
		out.push(boxLine(""));

		const title = " Annotate last message";
		const counter = `${this.annotations.length} note${this.annotations.length === 1 ? "" : "s"} `;
		const titlePad = Math.max(
			1,
			inner - visibleWidth(title) - visibleWidth(counter),
		);
		out.push(
			boxLine(
				theme.fg("accent", theme.bold(title)) +
					" ".repeat(titlePad) +
					theme.fg("muted", counter),
			),
		);

		for (let i = 0; i < bodyHeight; i++) {
			const rowIdx = this.scrollOffset + i;
			const row = rows[rowIdx];
			if (!row) {
				out.push(boxLine(""));
				continue;
			}
			const inSpan =
				span !== null && row.line >= span[0] && row.line <= span[1];
			const annIdx = this.annotationAt(row.line);
			const isCursor = row.line === this.cursor;

			if (inSpan) {
				const raw = `${isCursor ? ">" : " "}   ${row.text}`;
				out.push(borderV + theme.bg("selectedBg", padInner(raw)) + borderV);
				continue;
			}

			let badge = " ";
			if (annIdx >= 0) {
				const firstRow = rows.findIndex(
					(r) => r.line === this.annotations[annIdx]!.start,
				);
				badge =
					rowIdx === firstRow ? (annIdx < 9 ? String(annIdx + 1) : "+") : "│";
			}
			const gutter =
				(annIdx >= 0 ? theme.fg("success", badge) : badge) +
				(isCursor ? theme.fg("accent", ">") : " ") +
				"  ";
			out.push(boxLine(gutter + row.text));
		}

		if (this.mode === "note" && this.anchor !== null) {
			const s = Math.min(this.anchor, this.cursor);
			const e = Math.max(this.anchor, this.cursor);
			const range = s === e ? `line ${s + 1}` : `lines ${s + 1}-${e + 1}`;
			const label =
				this.editing !== null
					? ` edit note for ${range}:`
					: ` note for ${range}:`;
			out.push(boxLine(theme.fg("accent", label)));
			for (const line of this.noteInput.render(inner - 2)) {
				out.push(boxLine(` ${line}`));
			}
		}

		let help: string;
		if (this.hint) help = this.hint;
		else if (this.mode === "span")
			help = "j/k extend · s/enter note · esc cancel span";
		else if (this.mode === "note" && this.editing !== null)
			help = "enter save · empty note deletes · esc keep original";
		else if (this.mode === "note") help = "type note · enter save · esc cancel";
		else if (this.annotationAt(this.cursor) >= 0)
			help =
				"i/enter edit · x delete · C clear · v span · shift+enter/S send · esc";
		else
			help = "j/k move · v span · i/enter note · shift+enter/S send · esc";
		out.push(boxLine(` ${theme.fg("dim", help)}`));

		out.push(boxLine(""));
		out.push(theme.fg("accent", theme.bold(`┗${"━".repeat(inner)}┛`)));
		return out;
	}

	invalidate(): void {
		this.cachedWidth = undefined;
		this.rows = [];
		this.noteInput.invalidate();
	}

	private buildRows(width: number): Row[] {
		if (this.cachedWidth === width && this.rows.length > 0) return this.rows;
		const contentWidth = Math.max(8, width - 2 - GUTTER);
		const rows: Row[] = [];
		this.lines.forEach((line, i) => {
			for (const text of wrapLine(line, contentWidth)) {
				rows.push({ line: i, text });
			}
		});
		this.rows = rows;
		this.cachedWidth = width;
		return rows;
	}

	private moveCursor(delta: number): void {
		this.cursor = Math.max(
			0,
			Math.min(this.lines.length - 1, this.cursor + delta),
		);
		this.hint = null;
	}

	private spanRange(): [number, number] | null {
		if (this.anchor === null) return null;
		return [
			Math.min(this.anchor, this.cursor),
			Math.max(this.anchor, this.cursor),
		];
	}

	private annotationAt(line: number): number {
		for (let i = this.annotations.length - 1; i >= 0; i--) {
			const a = this.annotations[i]!;
			if (line >= a.start && line <= a.end) return i;
		}
		return -1;
	}

	private handleEscape(): void {
		if (this.mode === "span") {
			this.mode = "nav";
			this.anchor = null;
			return;
		}
		this.done(null);
	}

	private openNote(): void {
		this.mode = "note";
		this.noteInput.setValue("");
		this.hint = null;
	}

	private openEdit(annIdx: number): void {
		const a = this.annotations[annIdx]!;
		this.editing = annIdx;
		this.anchor = a.start;
		this.cursor = a.end;
		this.mode = "note";
		this.noteInput.setValue(a.note);
		// setValue clamps rather than moves the cursor; send ctrl+e so the user
		// appends to the existing note instead of overwriting from position 0.
		this.noteInput.handleInput("\x05");
		this.hint = null;
	}

	private commitNote(value: string): void {
		const note = value.trim();
		if (this.editing !== null) {
			if (note) {
				this.annotations[this.editing]!.note = note;
			} else {
				// Saving an emptied note removes the annotation entirely.
				this.annotations.splice(this.editing, 1);
			}
		} else if (note && this.anchor !== null) {
			this.annotations.push({
				start: Math.min(this.anchor, this.cursor),
				end: Math.max(this.anchor, this.cursor),
				note,
			});
		}
		this.resetNoteState();
	}

	private cancelNote(): void {
		this.resetNoteState();
	}

	private resetNoteState(): void {
		this.noteInput.setValue("");
		this.mode = "nav";
		this.anchor = null;
		this.editing = null;
	}
}

function lastAssistantTarget(
	ctx: ExtensionContext,
): { id: string; text: string } | null {
	const branch = ctx.sessionManager.getBranch();
	for (let i = branch.length - 1; i >= 0; i--) {
		const entry = branch[i]!;
		if (entry.type !== "message" || entry.message.role !== "assistant")
			continue;
		const content = entry.message.content;
		if (!Array.isArray(content)) continue;
		const text = content
			.filter((c): c is { type: "text"; text: string } => c.type === "text")
			.map((c) => c.text)
			.join("\n");
		if (text.trim()) return { id: entry.id, text };
	}
	return null;
}

export default function (pi: ExtensionAPI) {
	// Per-message draft state kept in memory for the session, so an annotation
	// session can span multiple overlay opens without re-marking anything.
	const drafts = new Map<string, Draft>();

	const launch = async (ctx: ExtensionContext) => {
		if (ctx.mode !== "tui") {
			ctx.ui.notify("annotate requires the TUI", "warning");
			return;
		}
		const target = lastAssistantTarget(ctx);
		if (!target) {
			ctx.ui.notify("no assistant message with text to annotate", "warning");
			return;
		}
		let overlay: AnnotateOverlay | undefined;
		const result = await ctx.ui.custom<Annotation[] | null>(
			(tui, theme, _keybindings, done) => {
				overlay = new AnnotateOverlay(
					tui,
					theme,
					target.text,
					done,
					drafts.get(target.id),
				);
				return overlay;
			},
			{
				overlay: true,
				overlayOptions: {
					width: "90%",
					minWidth: 60,
					maxHeight: "80%",
					anchor: "center",
				},
			},
		);
		if (!result || result.length === 0) {
			const state = overlay?.getDraftState();
			if (state && state.annotations.length > 0) {
				drafts.set(target.id, state);
				ctx.ui.notify(
					`draft saved - ${state.annotations.length} annotation(s); reopen with ctrl+shift+a`,
					"info",
				);
			} else {
				drafts.delete(target.id);
			}
			return;
		}
		drafts.delete(target.id);
		const message = formatAnnotations(result, target.text);
		if (ctx.isIdle()) {
			pi.sendUserMessage(message);
		} else {
			pi.sendUserMessage(message, { deliverAs: "followUp" });
		}
		ctx.ui.notify(
			`sent ${result.length} annotation${result.length === 1 ? "" : "s"}`,
			"info",
		);
	};

	pi.registerCommand("annotate", {
		description:
			"Annotate spans of the last assistant message and send them as one batch",
		handler: async (_args, ctx) => launch(ctx),
	});

	// Fires only on Kitty-protocol terminals (Ghostty, kitty, WezTerm); elsewhere
	// it arrives as plain ctrl+a and harmlessly does editor line-start instead.
	const shortcut = (process.env.PI_ANNOTATE_SHORTCUT ||
		"ctrl+shift+a") as KeyId;
	pi.registerShortcut(shortcut, {
		description: "Annotate the last assistant message",
		handler: async (ctx) => launch(ctx),
	});
}