/**
 * Ripgrep File Reference Extension
 *
 * Adds a `/r` command that uses a cached `rg --files` index for live command
 * argument completions. Submitting `/r <query>` inserts the selected file back
 * into the editor as an `@file` reference.
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { basename } from "node:path";

const MAX_COMPLETIONS = 50;
const RG_TIMEOUT_MS = 15000;

interface FileEntry {
	path: string;
	pathLower: string;
	fileNameLower: string;
	index: number;
}

interface RankedFile {
	entry: FileEntry;
	score: number;
}

function scorePath(queryLower: string, entry: FileEntry): number | null {
	if (entry.pathLower === queryLower) return 0;
	if (entry.fileNameLower === queryLower) return 1;
	if (entry.fileNameLower.startsWith(queryLower)) return 2;
	if (entry.pathLower.startsWith(queryLower)) return 3;
	if (entry.pathLower.includes(`/${queryLower}`)) return 4;
	if (entry.fileNameLower.includes(queryLower)) return 5;
	if (entry.pathLower.includes(queryLower)) return 6;
	return null;
}

function compareRankedFiles(a: RankedFile, b: RankedFile): number {
	if (a.score !== b.score) return a.score - b.score;
	if (a.entry.path.length !== b.entry.path.length) return a.entry.path.length - b.entry.path.length;
	return a.entry.index - b.entry.index;
}

function insertRankedFile(ranked: RankedFile[], candidate: RankedFile, limit: number): void {
	let insertIndex = ranked.length;
	for (let i = 0; i < ranked.length; i++) {
		if (compareRankedFiles(candidate, ranked[i]) < 0) {
			insertIndex = i;
			break;
		}
	}

	ranked.splice(insertIndex, 0, candidate);
	if (ranked.length > limit) {
		ranked.pop();
	}
}

export default function rgFileReferenceExtension(pi: ExtensionAPI) {
	let cwd = process.cwd();
	let fileCache: FileEntry[] = [];
	let refreshPromise: Promise<void> | null = null;
	let lastRefreshError: string | undefined;

	const refreshCache = async (nextCwd?: string): Promise<void> => {
		if (nextCwd && nextCwd !== cwd) {
			cwd = nextCwd;
			fileCache = [];
		}

		if (refreshPromise) {
			await refreshPromise;
			return;
		}

		refreshPromise = (async () => {
			const result = await pi.exec("rg", ["--files"], { cwd, timeout: RG_TIMEOUT_MS });
			if (result.killed) {
				lastRefreshError = "rg --files was interrupted";
				return;
			}
			if (result.code !== 0) {
				lastRefreshError = result.stderr.trim() || `rg --files exited with code ${result.code}`;
				return;
			}

			const files = result.stdout
				.split("\n")
				.map((line) => line.trim())
				.filter((line) => line.length > 0);

			fileCache = files.map((path, index) => ({
				path,
				pathLower: path.toLowerCase(),
				fileNameLower: basename(path).toLowerCase(),
				index,
			}));
			lastRefreshError = undefined;
		})();

		try {
			await refreshPromise;
		} finally {
			refreshPromise = null;
		}
	};

	const findExactMatch = (query: string): string | undefined => {
		return fileCache.find((entry) => entry.path === query)?.path;
	};

	const findMatches = (query: string, limit: number): string[] => {
		const queryLower = query.toLowerCase();
		const ranked: RankedFile[] = [];

		for (const entry of fileCache) {
			const score = scorePath(queryLower, entry);
			if (score === null) {
				continue;
			}
			insertRankedFile(ranked, { entry, score }, limit);
		}

		return ranked.map((item) => item.entry.path);
	};

	const insertReference = (path: string, setEditorText: (text: string) => void): void => {
		setEditorText(`@${path} `);
	};

	void refreshCache(cwd);

	pi.on("session_start", async (_event, ctx) => {
		await refreshCache(ctx.cwd);
	});

	pi.registerCommand("r", {
		description: "Insert an @file reference from a cached ripgrep file index",
		getArgumentCompletions: (prefix) => {
			const query = prefix.trim();
			if (!query || fileCache.length === 0) {
				return null;
			}

			const matches = findMatches(query, MAX_COMPLETIONS);
			return matches.length > 0 ? matches.map((path) => ({ value: path, label: path })) : null;
		},
		handler: async (args, ctx) => {
			const query = args.trim();
			if (!query) {
				ctx.ui.notify("Usage: /r <pattern>", "warning");
				return;
			}

			if (fileCache.length === 0) {
				await refreshCache(ctx.cwd);
			}

			let exactMatch = findExactMatch(query);
			let matches = exactMatch ? [exactMatch] : findMatches(query, MAX_COMPLETIONS);
			if (matches.length === 0) {
				await refreshCache(ctx.cwd);
				exactMatch = findExactMatch(query);
				matches = exactMatch ? [exactMatch] : findMatches(query, MAX_COMPLETIONS);
			}

			if (matches.length === 0) {
				ctx.ui.notify(lastRefreshError ? `No matches (${lastRefreshError})` : "No matches", "info");
				return;
			}

			const selected = matches.length === 1 ? matches[0] : await ctx.ui.select("Insert file reference", matches);
			if (!selected) {
				return;
			}

			insertReference(selected, ctx.ui.setEditorText.bind(ctx.ui));
		},
	});
}