import { CustomEditor, type ExtensionAPI, type KeybindingsManager } from "@earendil-works/pi-coding-agent";
import type { EditorTheme, TUI } from "@earendil-works/pi-tui";
import { matchesKey, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
import { messageText, printableInput, ReverseSearchController, type PromptItem } from "./reverse-search-lib/core.ts";

function fitSearchStatus(line: string, label: string, width: number): string {
	if (width <= 0) return "";
	if (visibleWidth(label) >= width) return truncateToWidth(label, width, "");

	const lineWidth = visibleWidth(line);
	if (lineWidth + visibleWidth(label) <= width) return line + label;

	return truncateToWidth(line, Math.max(0, width - visibleWidth(label)), "") + label;
}

const ANSI_SEQUENCE = /\x1b\[[0-?]*[ -/]*[@-~]|\x1b\][^\x07]*(?:\x07|\x1b\\)|\x1b[PX^_].*?\x1b\\/g;

function escapeRegExp(text: string): string {
	return text.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
}

function highlightSearchMatches(line: string, query: string): string {
	if (!query) return line;

	const queryRe = new RegExp(escapeRegExp(query), "gi");
	let out = "";
	let lastIndex = 0;

	for (const match of line.matchAll(ANSI_SEQUENCE)) {
		const index = match.index ?? 0;
		out += line.slice(lastIndex, index).replace(queryRe, (text) => `\x1b[7m${text}\x1b[27m`);
		out += match[0];
		lastIndex = index + match[0].length;
	}

	out += line.slice(lastIndex).replace(queryRe, (text) => `\x1b[7m${text}\x1b[27m`);
	return out;
}

class ReverseSearchEditor extends CustomEditor {
	private readonly search = new ReverseSearchController();

	constructor(
		tui: TUI,
		theme: EditorTheme,
		keybindings: KeybindingsManager,
		private readonly getPrompts: () => PromptItem[],
	) {
		super(tui, theme, keybindings);
	}

	handleInput(data: string): void {
		if (matchesKey(data, "ctrl+r")) {
			if (!this.search.active) {
				this.setText(this.search.start(this.getPrompts(), this.getText()));
			} else {
				this.setText(this.search.cycleOlder());
			}
			this.tui.requestRender();
			return;
		}

		if (!this.search.active) {
			super.handleInput(data);
			return;
		}

		if (matchesKey(data, "escape") || matchesKey(data, "ctrl+g")) {
			this.setText(this.search.cancel());
			this.tui.requestRender();
			return;
		}

		if (matchesKey(data, "ctrl+c")) {
			// Shell-like abort: cancel search, then let Pi's app-level Ctrl+C handling clear/interrupt.
			this.setText(this.search.cancel());
			super.handleInput(data);
			this.tui.requestRender();
			return;
		}

		if (matchesKey(data, "backspace")) {
			this.setText(this.search.backspaceQuery());
			this.tui.requestRender();
			return;
		}

		if (matchesKey(data, "delete")) {
			this.setText(this.search.clearQuery());
			this.tui.requestRender();
			return;
		}

		if (matchesKey(data, "up") || matchesKey(data, "ctrl+p")) {
			this.setText(this.search.cycleOlder());
			this.tui.requestRender();
			return;
		}

		if (matchesKey(data, "down") || matchesKey(data, "ctrl+n") || matchesKey(data, "ctrl+s")) {
			this.setText(this.search.cycleNewer());
			this.tui.requestRender();
			return;
		}

		const printable = printableInput(data);
		if (printable !== undefined) {
			this.setText(this.search.appendQuery(printable));
			this.tui.requestRender();
			return;
		}

		// Bash/readline and zsh leave incremental search when an editing/submission key is used.
		// Delegate the key so Enter submits immediately and arrows/editing operate on the chosen prompt.
		this.search.accept();
		super.handleInput(data);
		this.tui.requestRender();
	}

	render(width: number): string[] {
		const lines = super.render(width);
		if (this.search.active && lines.length > 0) {
			// Avoid transient single-character highlighting: Pi's incremental renderer can leave
			// stale style-only cells behind as the query grows from "i" to "idea".
			const query = this.search.statusMessage === "failed" || this.search.query.length < 2 ? "" : this.search.query;
			const highlighted = lines.map((line) => highlightSearchMatches(line, query));
			const last = highlighted.length - 1;
			highlighted[last] = fitSearchStatus(highlighted[last]!, this.search.searchLabel(), width);
			return highlighted;
		}
		return lines;
	}
}

function promptsFromEntries(entries: Array<Record<string, unknown>>): PromptItem[] {
	const prompts: PromptItem[] = [];

	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i] as { type?: unknown; message?: { role?: unknown; content?: unknown; timestamp?: unknown } };
		if (entry.type !== "message" || entry.message?.role !== "user") continue;

		const text = messageText(entry.message.content);
		if (!text) continue;
		prompts.push({
			text,
			timestamp: typeof entry.message.timestamp === "number" ? entry.message.timestamp : undefined,
		});
	}

	return prompts;
}

export default function (pi: ExtensionAPI) {
	pi.on("session_start", (_event, ctx) => {
		if (ctx.mode !== "tui") return;

		ctx.ui.setEditorComponent((tui, theme, keybindings) => {
			return new ReverseSearchEditor(tui, theme, keybindings, () => {
				return promptsFromEntries(ctx.sessionManager.getBranch() as Array<Record<string, unknown>>);
			});
		});
	});
}