import { readFileSync, existsSync } from "node:fs";
import { type Plugin, tool } from "@opencode-ai/plugin";

/**
 * Extension to language mapping for rule recommendations.
 * Only includes languages we have rule files for.
 */
const EXTENSION_TO_LANGUAGE: Record<string, string> = {
	// TypeScript/JavaScript
	ts: "typescript",
	tsx: "typescript",
	js: "javascript",
	jsx: "javascript",
	mjs: "javascript",
	cjs: "javascript",

	// Go
	go: "go",

	// Nix
	nix: "nix",

	// Rust
	rs: "rust",

	// Ruby
	rb: "ruby",
	erb: "ruby",
	rake: "ruby",

	// Python
	py: "python",
	pyw: "python",
	pyi: "python",

	// Java/C-family
	java: "java",
	c: "c",
	h: "c",
	cpp: "cpp",
	cc: "cpp",
	cxx: "cpp",
	hpp: "cpp",
	cs: "csharp",

	// Shell
	sh: "shell",
	bash: "shell",
	zsh: "shell",

	// Web
	html: "html",
	htm: "html",
	css: "css",
	scss: "css",
	sass: "css",
	less: "css",

	// Data/Config
	json: "json",
	yaml: "yaml",
	yml: "yaml",
	xml: "xml",
	toml: "toml",

	// Documentation
	md: "markdown",
	markdown: "markdown",

	// Database
	sql: "sql",
	ddl: "sql",
};

/**
 * Extract file extension and map to language name.
 */
function getLanguageFromFilePath(filePath: string): string | null {
	const match = filePath.match(/\.([a-zA-Z0-9]+)(?:\?.*)?$/);
	if (!match) return null;
	const ext = match[1]!.toLowerCase();
	return EXTENSION_TO_LANGUAGE[ext] ?? null;
}

/**
 * Registry to track which rules have been read in each session.
 */
const sessionRulesRead = new Map<string, Set<string>>();

/**
 * Cleanup old sessions to prevent memory leaks.
 * Should be called periodically or when sessions exceed a threshold.
 */
function cleanupOldSessions(maxSessions = 1000): void {
	if (sessionRulesRead.size > maxSessions) {
		const entries = Array.from(sessionRulesRead.entries());
		// Remove oldest half
		for (let i = 0; i < entries.length / 2; i++) {
			sessionRulesRead.delete(entries[i]![0]);
		}
	}
}

/**
 * Finds the path to a rule file for the given language.
 * Returns the path if found, null otherwise.
 */
function findRuleFile(directory: string, language: string): string | null {
	const paths = [`${directory}/.rules/${language}.md`];

	const rulesDir = process.env.LANGRULES_DIR;
	if (rulesDir) {
		paths.push(`${rulesDir}/${language}.md`);
	}

	for (const rulePath of paths) {
		if (existsSync(rulePath)) {
			return rulePath;
		}
	}
	return null;
}

/**
 * Checks if rules should be prompted for a session/language combo.
 * Returns true if this is the first time, and marks it as read.
 */
function shouldPromptForRules(sessionID: string, language: string): boolean {
	if (!sessionRulesRead.has(sessionID)) {
		sessionRulesRead.set(sessionID, new Set());
	}

	const rulesSet = sessionRulesRead.get(sessionID);
	if (!rulesSet || rulesSet.has(language)) {
		return false;
	}

	rulesSet.add(language);
	return true;
}

export const LangRulesPlugin: Plugin = async ({ client, directory }) => {
	return {
		tool: {
			langrules: tool({
				description:
					"Read the critical rules needed for style and necessary context",
				args: {
					language: tool.schema
						.string()
						.describe(
							"Specific rule file to read (e.g., 'typescript', 'go', 'javascript')",
						),
				},
				async execute(args, toolCtx) {
					const sessionID = toolCtx?.sessionID;
					const languageFile = args.language;

					if (sessionID) {
						if (!sessionRulesRead.has(sessionID)) {
							sessionRulesRead.set(sessionID, new Set());
							cleanupOldSessions();
						}

						const rulesSet = sessionRulesRead.get(sessionID);
						if (rulesSet) {
							rulesSet.add(languageFile);
						}
					}

					const rulePath = findRuleFile(directory, languageFile);
					if (!rulePath) {
						return `No rule file found for '${languageFile}'`;
					}

					try {
						const content = readFileSync(rulePath, "utf-8");
						if (content.trim()) return content.trim();
					} catch (error) {}

					return `No rule file found for '${languageFile}'`;
				},
			}),
		},

		"tool.execute.before": async (input, output) => {
			const { sessionID, tool: toolName } = input;
			const { args } = output;

			if (!sessionID) return;

			// Only prompt for edit and write tools to avoid context bloat during research
			if (toolName !== "edit" && toolName !== "write") return;

			const filePath = (args?.filePath as string) || (args?.path as string);
			if (!filePath || typeof filePath !== "string") return;

			const language = getLanguageFromFilePath(filePath);
			if (!language) return;

			if (!findRuleFile(directory, language)) return;

			if (shouldPromptForRules(sessionID, language)) {
				const ext = filePath.split(".").pop()?.toLowerCase() || "";
				client.session.prompt({
					path: {
						id: sessionID,
					},
					body: {
						parts: [
							{
								synthetic: true,
								type: "text",
								text: `**Hint:** You're about to modify a \`${ext}\` file (${language}). Consider reading the ${language} coding rules: \`langrules({ language: '${language}' })\``,
							},
						],
					},
				});
			}
		},
	};
};