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

/**
 * Pi Personas Extension
 *
 * Switch between agent personas stored as markdown files.
 *
 * Personas live in the Pi agent directory's personas/ subdirectory 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 as the "persona" key
 * in the Pi agent directory's settings.json.
 */

import * as fs from "node:fs";
import * as path from "node:path";
import { getAgentDir, type ExtensionAPI, SettingsManager } from "@earendil-works/pi-coding-agent";

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

type JsonValue = JsonObject | JsonValue[] | string | number | boolean | null;
type JsonObject = { [key: string]: JsonValue };

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

function getSettingsPath(): string {
	return path.join(getAgentDir(), "settings.json");
}

/** 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 [];
	}
}

function parseSettingsJson(content: string): JsonObject {
	if (!content.trim()) return {};
	const parsed: JsonValue = JSON.parse(content);
	if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
		throw new Error("settings.json must contain a JSON object");
	}
	return parsed;
}

function readPersonaSetting(cwd: string): string | null {
	const settingsManager = SettingsManager.create(cwd, getAgentDir());
	const settings = settingsManager.getGlobalSettings();
	const descriptor = Object.getOwnPropertyDescriptor(settings, ACTIVE_PERSONA_SETTING);
	if (!descriptor || !("value" in descriptor)) return null;
	if (typeof descriptor.value !== "string") return null;
	const name = descriptor.value.trim();
	return name ? name : null;
}

function writePersonaSetting(name: string | null): void {
	const filePath = getSettingsPath();
	try {
		const settings = fs.existsSync(filePath) ? parseSettingsJson(fs.readFileSync(filePath, "utf-8")) : {};

		if (name) {
			settings[ACTIVE_PERSONA_SETTING] = name;
		} else {
			delete settings[ACTIVE_PERSONA_SETTING];
		}

		const dir = path.dirname(filePath);
		if (!fs.existsSync(dir)) {
			fs.mkdirSync(dir, { recursive: true });
		}

		fs.writeFileSync(filePath, `${JSON.stringify(settings, null, 2)}\n`, "utf-8");
	} catch (err) {
		console.error(`Failed to write active persona setting: ${err}`);
	}
}

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

		// Reject names that could escape the personas directory
		const available = listPersonas();
		if (!available.includes(name)) {
			console.error(`Configured 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 setting to clear). */
function writeActivePersona(name: string | null): void {
	writePersonaSetting(name);
}

/** 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(ctx.cwd);
		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");
		},
	});
}