import { lstat, mkdir, readdir, readFile, writeFile } from "node:fs/promises";
import { homedir } from "node:os";
import { dirname, join, resolve, sep } from "node:path";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

export const DEFAULT_SYSTEMS_DIR = join(homedir(), ".pi", "agent", "systems");
export const DEFAULT_CONFIG_PATH = join(
	homedir(),
	".pi",
	"agent",
	"system-prompt-switcher.json",
);

const NO_SELECTION = "No system prompt selected.";

export interface SelectedPrompt {
	ok: true;
	name: string;
	path: string;
	content: string;
}

export interface PromptError {
	ok: false;
	error: string;
}

interface StoredConfig {
	active?: unknown;
}

function promptPath(systemsDir: string, name: string): string | undefined {
	if (!name || name.includes("/") || name.includes("\\")) return undefined;
	const base = resolve(systemsDir);
	const candidate = resolve(base, name);
	if (candidate !== base && candidate.startsWith(`${base}${sep}`))
		return candidate;
	return undefined;
}

export async function discoverSystemPrompts(
	systemsDir: string,
): Promise<string[]> {
	try {
		const entries = await readdir(systemsDir, { withFileTypes: true });
		return entries
			.filter((entry) => entry.isFile() && !entry.name.startsWith("."))
			.map((entry) => entry.name)
			.sort((left, right) => left.localeCompare(right));
	} catch (error) {
		if (error && typeof error === "object" && "code" in error) {
			const code = String(error.code);
			if (code === "ENOENT" || code === "ENOTDIR") return [];
		}
		throw error;
	}
}

export async function saveSelectedPrompt(
	configPath: string,
	name: string,
): Promise<void> {
	await mkdir(dirname(configPath), { recursive: true });
	await writeFile(
		configPath,
		`${JSON.stringify({ active: name }, null, "\t")}\n`,
		"utf8",
	);
}

export async function selectSystemPrompt(options: {
	configPath: string;
	systemsDir: string;
	name: string;
}): Promise<{ ok: true; name: string; path: string } | PromptError> {
	const name = options.name.trim();
	const path = promptPath(options.systemsDir, name);
	if (!path) return { ok: false, error: `Invalid system prompt name: ${name}` };

	try {
		const info = await lstat(path);
		if (!info.isFile()) {
			return { ok: false, error: `System prompt not found: ${name}` };
		}
	} catch (error) {
		if (error && typeof error === "object" && "code" in error) {
			const code = String(error.code);
			if (code === "ENOENT" || code === "ENOTDIR") {
				return { ok: false, error: `System prompt not found: ${name}` };
			}
		}
		const message = error instanceof Error ? error.message : String(error);
		return {
			ok: false,
			error: `Could not inspect system prompt ${name}: ${message}`,
		};
	}

	await saveSelectedPrompt(options.configPath, name);
	return { ok: true, name, path };
}

export async function loadSelectedPrompt(options: {
	configPath: string;
	systemsDir: string;
}): Promise<SelectedPrompt | PromptError> {
	let rawConfig: StoredConfig;
	try {
		rawConfig = JSON.parse(
			await readFile(options.configPath, "utf8"),
		) as StoredConfig;
	} catch (error) {
		if (error && typeof error === "object" && "code" in error) {
			const code = String(error.code);
			if (code === "ENOENT") return { ok: false, error: NO_SELECTION };
		}
		const message = error instanceof Error ? error.message : String(error);
		return {
			ok: false,
			error: `Could not read system prompt switcher config: ${message}`,
		};
	}

	if (typeof rawConfig.active !== "string" || !rawConfig.active.trim()) {
		return { ok: false, error: NO_SELECTION };
	}

	const name = rawConfig.active.trim();
	const path = promptPath(options.systemsDir, name);
	if (!path) return { ok: false, error: `Invalid system prompt name: ${name}` };

	try {
		const systemsInfo = await lstat(options.systemsDir);
		if (!systemsInfo.isDirectory()) return { ok: false, error: NO_SELECTION };
	} catch (error) {
		if (error && typeof error === "object" && "code" in error) {
			const code = String(error.code);
			if (code === "ENOENT" || code === "ENOTDIR") {
				return { ok: false, error: NO_SELECTION };
			}
		}
		const message = error instanceof Error ? error.message : String(error);
		return {
			ok: false,
			error: `Could not inspect systems directory ${options.systemsDir}: ${message}`,
		};
	}

	try {
		const info = await lstat(path);
		if (!info.isFile()) {
			return { ok: false, error: `System prompt not found: ${name}` };
		}
		return {
			ok: true,
			name,
			path,
			content: await readFile(path, "utf8"),
		};
	} catch (error) {
		const message = error instanceof Error ? error.message : String(error);
		return {
			ok: false,
			error: `Could not read selected system prompt ${name}: ${message}`,
		};
	}
}

export function appendActiveSystemPrompt(
	basePrompt: string,
	prompt: Pick<SelectedPrompt, "name" | "content">,
): string {
	return `${basePrompt}\n\n## Active system prompt: ${prompt.name}\n\n${prompt.content.trimEnd()}`;
}

export interface SystemPromptSwitcherOptions {
	configPath?: string;
	systemsDir?: string;
}

export default async function systemPromptSwitcherExtension(
	pi: ExtensionAPI,
	options: SystemPromptSwitcherOptions = {},
): Promise<void> {
	const configPath = options.configPath ?? DEFAULT_CONFIG_PATH;
	const systemsDir = options.systemsDir ?? DEFAULT_SYSTEMS_DIR;
	const activePrompt = await loadSelectedPrompt({ configPath, systemsDir });
	let warningShown = false;

	pi.on("session_start", (_event, ctx) => {
		if (!ctx.hasUI || !activePrompt.ok) return;
		ctx.ui.notify(`[System Prompt]\n  ${activePrompt.name}`, "info");
	});

	pi.on("before_agent_start", async (event, ctx) => {
		if (!activePrompt.ok) {
			if (activePrompt.error !== NO_SELECTION && !warningShown) {
				warningShown = true;
				ctx.ui.notify(activePrompt.error, "warning");
			}
			return undefined;
		}
		return {
			systemPrompt: appendActiveSystemPrompt(event.systemPrompt, activePrompt),
		};
	});

	pi.registerCommand("system-prompt", {
		description:
			"Switch active system prompt from ~/.pi/agent/systems/ and reload Pi",
		handler: async (args, ctx) => {
			const prompts = await discoverSystemPrompts(systemsDir);
			if (prompts.length === 0) {
				ctx.ui.notify(`No system prompts found in ${systemsDir}.`, "warning");
				return;
			}

			let name = args.trim();
			if (!name) {
				if (!ctx.hasUI) {
					ctx.ui.notify(`Pass one of: ${prompts.join(", ")}`, "warning");
					return;
				}
				const choice = await ctx.ui.select(
					"Select active system prompt:",
					prompts,
				);
				if (!choice) return;
				name = choice;
			}

			const result = await selectSystemPrompt({ configPath, systemsDir, name });
			if (!result.ok) {
				ctx.ui.notify(result.error, "error");
				return;
			}

			ctx.ui.notify(
				`Switched active system prompt to ${result.name}. Reloading Pi...`,
				"info",
			);
			await ctx.reload();
			return;
		},
	});
}