import { CustomEditor, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
import type {
	AutocompleteItem,
	AutocompleteProvider,
	AutocompleteSuggestions,
} from "@earendil-works/pi-tui";
import { matchesKey } from "@earendil-works/pi-tui";
import { existsSync } from "node:fs";
import { readFile } from "node:fs/promises";
import { createRequire } from "node:module";
import { homedir } from "node:os";
import { dirname, isAbsolute, join, normalize, resolve } from "node:path";
import { fileURLToPath } from "node:url";

export type Candidate = { text: string; source?: string };
export type FuzzyMatch = { text: string; source: string; score: number; exact: boolean };
export type FuzzyIndex = {
	addCandidatesAsync(items: Candidate[]): Promise<number>;
	queryAsync(needle: string, limit: number): Promise<FuzzyMatch[]>;
	len(): number;
};

type NativeModule = {
	FrizbeeIndex: new () => FuzzyIndex;
};

const MIN_PREFIX_LENGTH = 2;
const MAX_SUGGESTIONS = 8;
const QUERY_LIMIT = 16;
const ADD_BATCH_SIZE = 2_000;
const MAX_WORD_LENGTH = 64;
const MAX_PATH_LENGTH = 240;
const WORD_TRIGGER_CHARS = [
	..."abcdefghijklmnopqrstuvwxyz",
	..."ABCDEFGHIJKLMNOPQRSTUVWXYZ",
	..."0123456789",
	"_",
];

const require = createRequire(import.meta.url);
const extensionDir = dirname(fileURLToPath(import.meta.url));
const nativePath = join(
	extensionDir,
	"native-build",
	`${process.platform}-${process.arch}`,
	"frizbee_autocomplete.node",
);

function loadNative(): NativeModule | undefined {
	if (!existsSync(nativePath)) return undefined;
	try {
		return require(nativePath) as NativeModule;
	} catch {
		return undefined;
	}
}

function fuzzyScore(needle: string, candidate: string): number | undefined {
	const query = needle.toLowerCase();
	const text = candidate.toLowerCase();
	if (!query) return 0;
	if (text === query) return 100_000;

	let queryIndex = 0;
	let score = text.startsWith(query) ? 10_000 : 0;
	let previousMatch = -2;
	for (let textIndex = 0; textIndex < text.length && queryIndex < query.length; textIndex += 1) {
		if (text[textIndex] !== query[queryIndex]) continue;
		score += 100;
		if (textIndex === previousMatch + 1) score += 50;
		if (textIndex === 0 || /[-_\s/]/.test(text[textIndex - 1] ?? "")) score += 25;
		score -= Math.max(0, textIndex - previousMatch - 1);
		previousMatch = textIndex;
		queryIndex += 1;
	}

	if (queryIndex !== query.length) return undefined;
	return score - Math.max(0, text.length - query.length);
}

export class JavaScriptFuzzyIndex implements FuzzyIndex {
	private readonly candidates: Array<{ text: string; source: string }> = [];
	private readonly seen = new Set<string>();

	async addCandidatesAsync(items: Candidate[]): Promise<number> {
		let added = 0;
		for (const item of items) {
			const text = item.text.trim();
			if (!text) continue;
			const key = text.toLowerCase();
			if (this.seen.has(key)) continue;
			this.seen.add(key);
			this.candidates.push({ text, source: item.source ?? "unknown" });
			added += 1;
		}
		return added;
	}

	async queryAsync(needle: string, limit: number): Promise<FuzzyMatch[]> {
		if (!needle.trim() || limit <= 0) return [];
		return this.candidates
			.map((candidate) => ({ candidate, score: fuzzyScore(needle, candidate.text) }))
			.filter(
				(item): item is { candidate: { text: string; source: string }; score: number } =>
					item.score !== undefined,
			)
			.sort(
				(a, b) =>
					b.score - a.score ||
					a.candidate.text.length - b.candidate.text.length ||
					a.candidate.text.localeCompare(b.candidate.text),
			)
			.slice(0, limit)
			.map(({ candidate, score }) => ({
				...candidate,
				score,
				exact: candidate.text.toLowerCase() === needle.toLowerCase(),
			}));
	}

	len(): number {
		return this.candidates.length;
	}
}

function immediate(): Promise<void> {
	return new Promise((resolve) => setImmediate(resolve));
}

export function extractCurrentWordPrefix(textBeforeCursor: string): string | undefined {
	const match = textBeforeCursor.match(/[A-Za-z0-9_'-]+$/);
	const prefix = match?.[0];
	if (!prefix || prefix.length < MIN_PREFIX_LENGTH || /^\d+$/.test(prefix)) {
		return undefined;
	}
	return prefix;
}

export function shouldDelegateToBuiltIn(textBeforeCursor: string): boolean {
	const trimmedStart = textBeforeCursor.trimStart();
	if (trimmedStart.startsWith("/")) {
		return true;
	}

	// Keep pi's @ file completion and common extension trigger contexts intact.
	if (/(?:^|[\s])[@#][^\s]*$/.test(textBeforeCursor)) {
		return true;
	}

	// Do not steal path-like tokens from pi's existing path completion.
	const token = textBeforeCursor.match(/\S+$/)?.[0] ?? "";
	return token.includes("/") || token.startsWith(".") || token.startsWith("~");
}

export function replaceCurrentWordPrefix(
	lines: string[],
	cursorLine: number,
	cursorCol: number,
	value: string,
): { lines: string[]; cursorLine: number; cursorCol: number } {
	const nextLines = [...lines];
	const line = nextLines[cursorLine] ?? "";
	const before = line.slice(0, cursorCol);
	const after = line.slice(cursorCol);
	const prefixMatch = before.match(/[A-Za-z0-9_'-]+$/);
	if (!prefixMatch) {
		return { lines, cursorLine, cursorCol };
	}

	const startCol = cursorCol - prefixMatch[0].length;
	const suffixMatch = after.match(/^[A-Za-z0-9_'-]*/);
	const suffixLength = suffixMatch?.[0].length ?? 0;
	nextLines[cursorLine] = line.slice(0, startCol) + value + after.slice(suffixLength);
	return {
		lines: nextLines,
		cursorLine,
		cursorCol: startCol + value.length,
	};
}

function textFromMessage(message: any): string {
	const content = message?.content;
	if (typeof content === "string") {
		return content;
	}
	if (!Array.isArray(content)) {
		return "";
	}

	return content
		.map((part) => {
			if (!part || typeof part !== "object") return "";
			if (part.type === "text" && typeof part.text === "string") return part.text;
			if (typeof part.text === "string") return part.text;
			if (typeof part.content === "string") return part.content;
			return "";
		})
		.join("\n");
}

function sourceFromMessage(message: any): string {
	const role = typeof message?.role === "string" ? message.role : "message";
	if (role === "toolResult" && typeof message?.toolName === "string") {
		return `session:${message.toolName}`;
	}
	return `session:${role}`;
}

export function extractWords(text: string, source: string): Candidate[] {
	const candidates: Candidate[] = [];
	const matches = text.matchAll(/[A-Za-z][A-Za-z0-9_'-]{1,63}/g);
	for (const match of matches) {
		const word = match[0];
		if (word.length > MAX_WORD_LENGTH) {
			continue;
		}
		candidates.push({ text: word, source });
	}
	return candidates;
}

function resolveExistingPath(rawPath: string, cwd: string): string | undefined {
	let value = rawPath.trim().replace(/[),.;:!?]+$/g, "");
	if (!value || value.length > MAX_PATH_LENGTH || value.includes("://")) {
		return undefined;
	}

	if (value.startsWith("~/")) {
		value = join(homedir(), value.slice(2));
	}

	const absolutePath = normalize(isAbsolute(value) ? value : resolve(cwd, value));
	if (!existsSync(absolutePath)) {
		return undefined;
	}
	return absolutePath;
}

export function extractExistingPaths(text: string, cwd: string): Candidate[] {
	const candidates: Candidate[] = [];
	const seen = new Set<string>();
	const matches = text.matchAll(/(?:^|[\s(["'`])((?:~\/|\.{1,2}\/|\/|[A-Za-z0-9_.-]+\/)[A-Za-z0-9_@%+=:,./~#-]*[A-Za-z0-9_@%+=/~#-])/g);

	for (const match of matches) {
		const rawPath = match[1];
		if (!rawPath) {
			continue;
		}

		const fullPath = resolveExistingPath(rawPath, cwd);
		if (!fullPath || seen.has(fullPath)) {
			continue;
		}

		seen.add(fullPath);
		candidates.push({ text: fullPath, source: "session:path" });
	}

	return candidates;
}

function extractSessionCandidates(text: string, source: string, cwd: string): Candidate[] {
	return [...extractWords(text, source), ...extractExistingPaths(text, cwd)];
}

async function readDictionaryWords(): Promise<Candidate[]> {
	const dictionaryPath = "/usr/share/dict/words";
	if (!existsSync(dictionaryPath)) {
		return [];
	}

	const text = await readFile(dictionaryPath, "utf8");
	return text
		.split(/\r?\n/)
		.map((word) => word.trim())
		.filter((word) => /^[A-Za-z][A-Za-z0-9_'-]{1,63}$/.test(word))
		.map((word) => ({ text: word, source: "dictionary" }));
}

async function addCandidatesBatched(index: FuzzyIndex, candidates: Candidate[]): Promise<number> {
	let added = 0;
	for (let i = 0; i < candidates.length; i += ADD_BATCH_SIZE) {
		const batch = candidates.slice(i, i + ADD_BATCH_SIZE);
		added += await index.addCandidatesAsync(batch);
		await immediate();
	}
	return added;
}

function stripAnsiSequences(text: string): string {
	return text.replace(/\x1b\[[0-?]*[ -/]*[@-~]/g, "");
}

class WordAutocompleteEditor extends CustomEditor {
	private hasAutocompleteProvider = false;

	setAutocompleteProvider(provider: AutocompleteProvider): void {
		this.hasAutocompleteProvider = true;
		super.setAutocompleteProvider(provider);
	}

	render(width: number): string[] {
		const lines = super.render(width);
		if (!this.hasAutocompleteProvider) return lines;

		const autocompleteList = (this as unknown as {
			autocompleteList?: { render(listWidth: number): string[] };
		}).autocompleteList;
		const paddingX = this.getPaddingX();
		const contentWidth = Math.max(1, width - paddingX * 2);
		const listHeight =
			this.isShowingAutocomplete() && autocompleteList ? autocompleteList.render(contentWidth).length : 0;
		const reservedHeight = this.getAutocompleteMaxVisible() + 1;

		// Newer TUI versions reserve this space themselves. Detect that by the
		// trailing blank rows; otherwise pad the inherited editor output here.
		const lastLine = lines.at(-1);
		if (!lastLine || stripAnsiSequences(lastLine).trim().length === 0) return lines;

		const extraRows = Math.max(0, reservedHeight - listHeight);
		return extraRows === 0
			? lines
			: [...lines, ...Array.from({ length: extraRows }, () => " ".repeat(Math.max(0, width)))];
	}

	private isShowingSlashCommandAutocomplete(): boolean {
		const prefix = (this as unknown as { autocompletePrefix?: unknown }).autocompletePrefix;
		return this.isShowingAutocomplete() && typeof prefix === "string" && prefix.startsWith("/");
	}

	private cycleAutocompleteSelection(direction: 1 | -1): boolean {
		const autocompleteList = (this as unknown as {
			autocompleteList?: {
				filteredItems?: unknown[];
				selectedIndex?: number;
				setSelectedIndex?: (index: number) => void;
			};
		}).autocompleteList;
		const items = autocompleteList?.filteredItems;
		const selectedIndex = autocompleteList?.selectedIndex;
		if (!autocompleteList || !Array.isArray(items) || items.length === 0 || typeof selectedIndex !== "number") {
			return false;
		}

		const nextIndex = direction === 1
			? (selectedIndex + 1) % items.length
			: (selectedIndex - 1 + items.length) % items.length;
		autocompleteList.setSelectedIndex?.(nextIndex);
		(this as unknown as { tui?: { requestRender?: () => void } }).tui?.requestRender?.();
		return true;
	}

	handleInput(data: string): void {
		if (this.isShowingAutocomplete()) {
			if (this.isShowingSlashCommandAutocomplete()) {
				if (matchesKey(data, "enter")) {
					// Pi normally accepts slash completion and then submits on Enter.
					// Treat Enter as Tab so the command is inserted for further editing.
					super.handleInput("\t");
					return;
				}
			} else {
				if (matchesKey(data, "tab") && this.cycleAutocompleteSelection(1)) {
					return;
				}
				if (matchesKey(data, "shift+tab") && this.cycleAutocompleteSelection(-1)) {
					return;
				}
			}
		}

		const beforeText = this.getText();
		const beforeCursor = this.getCursor();
		const wasShowingAutocomplete = this.isShowingAutocomplete();
		super.handleInput(data);

		if (wasShowingAutocomplete || this.isShowingAutocomplete()) {
			return;
		}

		const afterText = this.getText();
		const afterCursor = this.getCursor();
		if (beforeText === afterText && beforeCursor.line === afterCursor.line && beforeCursor.col === afterCursor.col) {
			return;
		}

		const line = this.getLines()[afterCursor.line] ?? "";
		const textBeforeCursor = line.slice(0, afterCursor.col);
		if (shouldDelegateToBuiltIn(textBeforeCursor) || !extractCurrentWordPrefix(textBeforeCursor)) {
			return;
		}

		// tryTriggerAutocomplete is private in TypeScript, but public at runtime.
		// This keeps the existing pi autocomplete provider/list UI and only fixes
		// word-trigger timing for normal typing.
		(this as unknown as { tryTriggerAutocomplete?: () => void }).tryTriggerAutocomplete?.();
	}
}

function createProvider(index: FuzzyIndex, current: AutocompleteProvider): AutocompleteProvider {
	let latestQueryId = 0;

	return {
		triggerCharacters: WORD_TRIGGER_CHARS,

		async getSuggestions(lines, cursorLine, cursorCol, options): Promise<AutocompleteSuggestions | null> {
			const line = lines[cursorLine] ?? "";
			const beforeCursor = line.slice(0, cursorCol);

			if (shouldDelegateToBuiltIn(beforeCursor)) {
				return current.getSuggestions(lines, cursorLine, cursorCol, options);
			}

			const prefix = extractCurrentWordPrefix(beforeCursor);
			if (!prefix) {
				return options.force ? current.getSuggestions(lines, cursorLine, cursorCol, options) : null;
			}

			const queryId = ++latestQueryId;
			let matches: FuzzyMatch[];
			try {
				matches = await index.queryAsync(prefix, QUERY_LIMIT);
			} catch {
				return null;
			}
			if (options.signal.aborted || queryId !== latestQueryId) {
				return null;
			}

			const items: AutocompleteItem[] = matches
				.filter((match) => match.text.toLowerCase() !== prefix.toLowerCase())
				.slice(0, MAX_SUGGESTIONS)
				.map((match) => ({
					value: match.text,
					label: match.text,
					description: `${match.source} · score ${match.score}`,
				}));

			if (items.length === 0) {
				return null;
			}

			return { prefix, items };
		},

		applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
			const line = lines[cursorLine] ?? "";
			const beforeCursor = line.slice(0, cursorCol);
			if (shouldDelegateToBuiltIn(beforeCursor)) {
				return current.applyCompletion(lines, cursorLine, cursorCol, item, prefix);
			}
			return replaceCurrentWordPrefix(lines, cursorLine, cursorCol, item.value);
		},

		shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
			return current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ?? true;
		},
	};
}

export default function (pi: ExtensionAPI): void {
	const native = loadNative();
	const index: FuzzyIndex = native ? new native.FrizbeeIndex() : new JavaScriptFuzzyIndex();
	let addQueue: Promise<void> = Promise.resolve();

	const enqueueCandidates = (candidates: Candidate[], ctx?: { ui?: { notify?: (message: string, level: "error") => void } }) => {
		if (candidates.length === 0) {
			return;
		}

		addQueue = addQueue
			.then(async () => {
				await addCandidatesBatched(index, candidates);
			})
			.catch((error) => {
				ctx?.ui?.notify?.(`frizbee-autocomplete: indexing failed: ${error.message ?? String(error)}`, "error");
			});
	};

	pi.on("session_start", async (_event, ctx) => {
		if (!ctx.ui.getEditorComponent()) {
			ctx.ui.setEditorComponent((tui, theme, keybindings) => new WordAutocompleteEditor(tui, theme, keybindings));
		}

		for (const entry of ctx.sessionManager.getBranch()) {
			if (entry.type !== "message") {
				continue;
			}
			const text = textFromMessage(entry.message);
			enqueueCandidates(extractSessionCandidates(text, sourceFromMessage(entry.message), ctx.cwd), ctx);
		}

		if (native) {
			void readDictionaryWords()
				.then((words) => enqueueCandidates(words, ctx))
				.catch((error) => {
					ctx.ui.notify(`frizbee-autocomplete: failed to read dictionary: ${error.message ?? String(error)}`, "error");
				});
		}

		ctx.ui.addAutocompleteProvider((current) => createProvider(index, current));
	});

	pi.on("message_end", async (event, ctx) => {
		const text = textFromMessage(event.message);
		enqueueCandidates(extractSessionCandidates(text, sourceFromMessage(event.message), ctx.cwd), ctx);
	});
}