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

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

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 getPersonalitiesDir(): string {
	return path.join(getAgentDir(), "personalities");
}

function getActivePersonalityPath(): string {
	return path.join(getAgentDir(), "personality");
}

/** List available personality names (without .md extension). */
function listPersonalities(): string[] {
	const dir = getPersonalitiesDir();
	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 personalities: ${err}`);
		return [];
	}
}

/**
 * Read the persisted active personality name, or null if none.
 * Validates the name against available personalities to prevent
 * path traversal via a tampered persistence file.
 */
function readActivePersonality(): string | null {
	const filePath = getActivePersonalityPath();
	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 personalities directory
		const available = listPersonalities();
		if (!available.includes(name)) {
			console.error(`Persisted personality "${name}" not in available list, ignoring`);
			return null;
		}

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

/** Persist the active personality name (or remove the file to clear). */
function writeActivePersonality(name: string | null): void {
	const filePath = getActivePersonalityPath();
	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 personality: ${err}`);
	}
}

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

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

function wrapPersonality(content: string): string {
	return `\n\n${PERSONALITY_TAG_OPEN}\n${content}\n${PERSONALITY_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 activePersonality: string | null = null;
	let sessionPersonalityContent: string | null = null;
	let pendingReminder: string | null = null;

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

	// Lock personality 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 = readActivePersonality();
		pendingReminder = null;
		sessionPersonalityContent = null;

		if (persisted) {
			const content = readPersonalityContent(persisted);
			if (content) {
				activePersonality = persisted;
				sessionPersonalityContent = content;
				ctx.ui.notify(`Set personality to "${persisted}"`, "info");
			} else {
				// Personality file was removed since last session
				setActive(null);
			}
		}
	});

	// Append the session personality 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 (!sessionPersonalityContent && !pendingReminder) return undefined;

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

		if (sessionPersonalityContent) {
			result.systemPrompt = event.systemPrompt + wrapPersonality(sessionPersonalityContent);
		}

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

		return result;
	});

	// Register /personality command
	pi.registerCommand("personality", {
		description: "Switch agent personality or clear it",
		getArgumentCompletions: (prefix: string) => {
			const names = [...listPersonalities(), ...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() || "";

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

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

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

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

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

			if (!choice) return; // cancelled

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

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

			const content = readPersonalityContent(picked);
			if (!content) {
				ctx.ui.notify(`Personality "${picked}" could not be read`, "error");
				return;
			}
			pendingReminder = wrapPersonality(content);
			setActive(picked);
			ctx.ui.notify(`Set personality to "${picked}"`, "info");
		},
	});
}