import type { BuildSystemPromptOptions, ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { mkdir, writeFile } from "node:fs/promises";
import { dirname, resolve } from "node:path";

const ROUTING_MARKER = "<!-- my-pi-tool-routing -->";
const DUMP_RELATIVE_PATH = ".pi-debug/effective-system-prompt.md";

type SelectedTools = BuildSystemPromptOptions["selectedTools"];

function hasTool(selectedTools: SelectedTools, name: string): boolean {
	return selectedTools === undefined || selectedTools.includes(name);
}

export function buildToolRoutingSection(selectedTools: SelectedTools): string {
	const rules: string[] = [];

	if (hasTool(selectedTools, "codegraph_explore")) {
		rules.push(
			"- Use codegraph_explore first for architecture, symbol relationships, code flows, and symbols you are about to change. Treat source returned by it as already read; do not repeat the same exploration with read or search unless information is missing.",
		);
	}

	if (
		hasTool(selectedTools, "lsp_definition") ||
		hasTool(selectedTools, "lsp_references") ||
		hasTool(selectedTools, "lsp_diagnostics")
	) {
		rules.push(
			"- Use LSP for symbol definitions, references, hover/type information, and post-edit diagnostics when the language server supports the file.",
		);
	}

	if (hasTool(selectedTools, "find") || hasTool(selectedTools, "grep") || hasTool(selectedTools, "multi_grep")) {
		rules.push(
			"- For literal search, narrow in stages instead of requesting a repository-wide dump:",
			"  1. Use find to identify likely files or directories. Constrain path/glob and keep the result limit small.",
			"  2. Use grep or multi_grep only inside the narrowed path or file to obtain matching line numbers. Start with a specific or literal pattern, a small context window, and a bounded limit.",
			"  3. Use read with offset/limit only for the exact matching region when full source is still needed.",
			"- Do not run a broad search expected to return hundreds or thousands of files or matches. If results truncate or hit a limit, narrow the path, glob, or pattern instead of increasing the limit or dumping all results.",
		);
	}

	if (hasTool(selectedTools, "ctx_execute") || hasTool(selectedTools, "ctx_execute_file")) {
		rules.push(
			"- Use Context Mode for logs, test/build output, generated data, large-file exploration/analysis/summarization, and any command whose output may be large or unpredictable. Prefer ctx_execute_file over read when deriving an answer from a large file inside the workspace, and print only the derived answer needed for the task.",
		);
	}

	if (hasTool(selectedTools, "ctx_execute_file")) {
		rules.push(
			"- ctx_execute_file is confined to the current project root. Do not call it for absolute paths outside the workspace, ../ traversal, or symlinks that resolve outside the workspace; host permission approval does not bypass this Context Mode boundary. Use another explicitly authorized tool for those files.",
		);
	}

	if (hasTool(selectedTools, "read")) {
		rules.push(
			"- Use read only when exact source is needed or the file is small. Before editing, request the smallest useful offset/limit range; do not read a complete large file merely to copy, compare, hash, count, explore, or summarize it.",
		);
	}

	if (hasTool(selectedTools, "read") && hasTool(selectedTools, "edit")) {
		rules.push(
			"- For existing-file changes, obtain fresh LINE#HASH anchors with a small read, then use anchored edit operations. Reuse fresh anchors returned by successful edits for chained changes; re-read only when anchors are stale or the next target region was not shown.",
		);
	}

	rules.push(
		"- Stop searching once the current evidence is sufficient. Avoid repeating the same lookup through CodeGraph, find, grep, and read without a concrete information gap.",
	);

	return `${ROUTING_MARKER}\n## Tool Routing\n\n${rules.join("\n")}`;
}

export function appendToolRouting(systemPrompt: string, selectedTools: SelectedTools): string {
	if (systemPrompt.includes(ROUTING_MARKER)) {
		return systemPrompt;
	}
	return `${systemPrompt}\n\n${buildToolRoutingSection(selectedTools)}`;
}

export function getSystemPromptDumpPath(cwd: string): string {
	return resolve(cwd, DUMP_RELATIVE_PATH);
}

export default function toolRoutingExtension(pi: ExtensionAPI): void {
	let lastEffectivePrompt: string | undefined;

	pi.on("before_agent_start", (event) => {
		lastEffectivePrompt = appendToolRouting(event.systemPrompt, event.systemPromptOptions.selectedTools);
		return { systemPrompt: lastEffectivePrompt };
	});

	pi.registerCommand("dump-system-prompt", {
		description: "Write the effective Pi system prompt to .pi-debug/effective-system-prompt.md",
		handler: async (_args, ctx) => {
			const outputPath = getSystemPromptDumpPath(ctx.cwd);
			const selectedTools = ctx.getSystemPromptOptions().selectedTools;
			const effectivePrompt = lastEffectivePrompt ?? appendToolRouting(ctx.getSystemPrompt(), selectedTools);

			await mkdir(dirname(outputPath), { recursive: true });
			await writeFile(outputPath, effectivePrompt, "utf8");
			ctx.ui.notify(`System prompt written to ${outputPath} (${effectivePrompt.length} chars)`, "info");
		},
	});
}