// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
//
// SPDX-License-Identifier: Unlicense

/**
 * Pi Personas Extension
 *
 * Switch between agent personas stored as markdown files.
 *
 * Personas live in $PI_CODING_AGENT_DIR/personas/ as .md files.
 * The active persona is appended to the system prompt in a <persona>
 * section, independent of SYSTEM.md and AGENTS.md.
 *
 * Commands:
 *   /persona              — Interactive picker to switch persona
 *   /persona none|unset|clear — Remove active persona
 *
 * The active persona persists across sessions in a plain text file
 * at $PI_CODING_AGENT_DIR/persona.
 */

import * as fs from "node:fs";
import * as path from "node:path";
import * as os from "node:os";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

const CLEAR_KEYWORDS = ["none", "unset", "clear"];

function getAgentDir(): string {
	const envDir = process.env.PI_CODING_AGENT_DIR;
	if (envDir) {
		if (envDir === "~") return os.homedir();
		if (envDir.startsWith("~/")) return path.join(os.homedir(), envDir.slice(2));
		return envDir;
	}
	return path.join(os.homedir(), ".pi", "agent");
}

function getPersonasDir(): string {
	return path.join(getAgentDir(), "personas");
}

function getActivePersonaPath(): string {
	return path.join(getAgentDir(), "persona");
}

/** List available persona names (without .md extension). */
function listPersonas(): string[] {
	const dir = getPersonasDir();
	if (!fs.existsSync(dir)) return [];

	try {
		return fs
			.readdirSync(dir)
			.filter((f) => f.endsWith(".md"))
			.map((f) => f.slice(0, -3))
			.sort();
	} catch (err) {
		console.error(`Failed to list personas: ${err}`);
		return [];
	}
}

/**
 * Read the persisted active persona name, or null if none.
 * Validates the name against available personas to prevent
 * path traversal via a tampered persistence file.
 */
function readActivePersona(): string | null {
	const filePath = getActivePersonaPath();
	if (!fs.existsSync(filePath)) return null;

	try {
		const name = fs.readFileSync(filePath, "utf-8").trim();
		if (!name) return null;

		// Reject names that could escape the personas directory
		const available = listPersonas();
		if (!available.includes(name)) {
			console.error(`Persisted persona "${name}" not in available list, ignoring`);
			return null;
		}

		return name;
	} catch (err) {
		console.error(`Failed to read active persona: ${err}`);
		return null;
	}
}

/** Persist the active persona name (or remove the file to clear). */
function writeActivePersona(name: string | null): void {
	const filePath = getActivePersonaPath();
	try {
		if (name) {
			fs.writeFileSync(filePath, `${name}\n`, "utf-8");
		} else if (fs.existsSync(filePath)) {
			fs.unlinkSync(filePath);
		}
	} catch (err) {
		console.error(`Failed to write active persona: ${err}`);
	}
}

/** Read persona markdown content, or null if file missing. */
function readPersonaContent(name: string): string | null {
	const filePath = path.join(getPersonasDir(), `${name}.md`);
	if (!fs.existsSync(filePath)) return null;
	try {
		return fs.readFileSync(filePath, "utf-8");
	} catch (err) {
		console.error(`Failed to read persona "${name}": ${err}`);
		return null;
	}
}

const PERSONA_TAG_OPEN =
	'<persona section_description="This section defines your personality. Fully embody the character, voice, and behaviour described within.">';
const PERSONA_TAG_CLOSE = "</persona>";

function wrapPersona(content: string): string {
	return `\n\n${PERSONA_TAG_OPEN}\n${content}\n${PERSONA_TAG_CLOSE}\n`;
}

export default function (pi: ExtensionAPI) {
	// "active" tracks what the user last picked (persisted to disk).
	// "session" is locked at session_start and never changes — it
	// controls what goes into the system prompt for this session.
	let activePersona: string | null = null;
	let sessionPersonaContent: string | null = null;
	let pendingReminder: string | null = null;

	/** Set the active persona name. Does NOT touch session prompt. */
	function setActive(name: string | null): void {
		activePersona = name;
		writeActivePersona(name);
	}

	// Lock persona into the system prompt at session start.
	// This content never changes for the lifetime of the session.
	pi.on("session_start", async (_event, ctx) => {
		const persisted = readActivePersona();
		pendingReminder = null;
		sessionPersonaContent = null;

		if (persisted) {
			const content = readPersonaContent(persisted);
			if (content) {
				activePersona = persisted;
				sessionPersonaContent = content;
				ctx.ui.notify(`Set persona to "${persisted}"`, "info");
			} else {
				// Persona file was removed since last session
				setActive(null);
			}
		}
	});

	// Append the session persona to the system prompt (ephemeral,
	// never persisted). The content is identical every turn so the
	// provider cache stays valid.
	//
	// Mid-session switches only fire a one-shot message — they never
	// alter the system prompt.
	pi.on("before_agent_start", async (event) => {
		if (!sessionPersonaContent && !pendingReminder) return undefined;

		const result: {
			systemPrompt?: string;
			message?: { customType: string; content: string; display: boolean };
		} = {};

		if (sessionPersonaContent) {
			result.systemPrompt = event.systemPrompt + wrapPersona(sessionPersonaContent);
		}

		if (pendingReminder) {
			result.message = {
				customType: "persona-switch",
				content: pendingReminder,
				display: false,
			};
			pendingReminder = null;
		}

		return result;
	});

	// Register /persona command
	pi.registerCommand("persona", {
		description: "Switch agent persona or clear it",
		getArgumentCompletions: (prefix: string) => {
			const names = [...listPersonas(), ...CLEAR_KEYWORDS];
			const filtered = names.filter((n) => n.startsWith(prefix)).map((n) => ({ value: n, label: n }));
			return filtered.length > 0 ? filtered : null;
		},
		handler: async (args, ctx) => {
			const arg = args?.trim() || "";

			// /persona none|unset|clear — remove active persona
			if (CLEAR_KEYWORDS.includes(arg.toLowerCase())) {
				if (!activePersona) {
					ctx.ui.notify("No persona is active", "info");
					return;
				}
				const was = activePersona;
				pendingReminder = `Persona "${was}" has been cleared. Revert to your default behaviour.`;
				setActive(null);
				ctx.ui.notify(`Persona cleared (was: ${was})`, "info");
				return;
			}

			// /persona someName — direct switch (not in picker)
			if (arg) {
				const available = listPersonas();
				if (!available.includes(arg)) {
					ctx.ui.notify(`Persona "${arg}" not found`, "error");
					return;
				}
				const content = readPersonaContent(arg);
				if (!content) {
					ctx.ui.notify(`Persona "${arg}" could not be read`, "error");
					return;
				}
				pendingReminder = wrapPersona(content);
				setActive(arg);
				ctx.ui.notify(`Set persona to "${arg}"`, "info");
				return;
			}

			// /persona — interactive picker
			const available = listPersonas();
			if (available.length === 0) {
				ctx.ui.notify("No personas available", "info");
				return;
			}

			const options = available.map((name) => (name === activePersona ? `${name} (active)` : name));
			const unsetLabel = activePersona ? "unset" : "unset (active)";
			options.unshift(unsetLabel);

			const choice = await ctx.ui.select("Pick a persona:", options);

			if (!choice) return; // cancelled

			// Strip " (active)" suffix to get the real name
			const picked = choice.replace(/ \(active\)$/, "");

			if (picked === "unset") {
				if (!activePersona) {
					ctx.ui.notify("No persona is active", "info");
				} else {
					const was = activePersona;
					pendingReminder = `Persona "${was}" has been cleared. Revert to your default behaviour.`;
					setActive(null);
					ctx.ui.notify(`Persona cleared (was: ${was})`, "info");
				}
				return;
			}

			const content = readPersonaContent(picked);
			if (!content) {
				ctx.ui.notify(`Persona "${picked}" could not be read`, "error");
				return;
			}
			pendingReminder = wrapPersona(content);
			setActive(picked);
			ctx.ui.notify(`Set persona to "${picked}"`, "info");
		},
	});
}