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

	return fs
		.readdirSync(dir)
		.filter((f) => f.endsWith(".md"))
		.map((f) => f.slice(0, -3))
		.sort();
}

/** Read the persisted active personality name, or null if none. */
function readActivePersonality(): string | null {
	const filePath = getActivePersonalityPath();
	if (!fs.existsSync(filePath)) return null;

	const name = fs.readFileSync(filePath, "utf-8").trim();
	return name || null;
}

/** Persist the active personality name (or remove the file to clear). */
function writeActivePersonality(name: string | null): void {
	const filePath = getActivePersonalityPath();
	if (name) {
		fs.writeFileSync(filePath, `${name}\n`, "utf-8");
	} else if (fs.existsSync(filePath)) {
		fs.unlinkSync(filePath);
	}
}

/** 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;
	return fs.readFileSync(filePath, "utf-8");
}

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) {
	let activePersonality: string | null = null;
	let personalityContent: string | null = null;
	let isFirstTurn = true;

	/** Load personality by name. Returns true if found and loaded. */
	function loadPersonality(name: string): boolean {
		const content = readPersonalityContent(name);
		if (!content) return false;

		activePersonality = name;
		personalityContent = content;
		return true;
	}

	/** Clear active personality (both in-memory and on disk). */
	function clearPersonality(): void {
		activePersonality = null;
		personalityContent = null;
		writeActivePersonality(null);
	}

	/** Switch to a personality: load into memory and persist to disk. */
	function switchPersonality(name: string): boolean {
		if (!loadPersonality(name)) return false;
		writeActivePersonality(name);
		return true;
	}

	// Restore persisted personality on session start
	pi.on("session_start", async (_event, ctx) => {
		const persisted = readActivePersonality();

		// Determine whether this is a fresh session or a resumed one
		const entries = ctx.sessionManager.getEntries();
		const hasMessages = entries.some((e) => e.type === "message" && e.message.role === "assistant");
		isFirstTurn = !hasMessages;

		if (persisted) {
			if (loadPersonality(persisted)) {
				ctx.ui.notify(`Set personality to "${persisted}"`, "info");
			} else {
				// Personality file was removed since last session
				clearPersonality();
			}
		}
	});

	// Inject personality into the system prompt or as a reminder message.
	// On the first turn (no prior messages), we modify the system prompt
	// directly so the personality is baked into the initial cache.
	// On subsequent turns, we inject a reminder message instead to avoid
	// invalidating the inference cache.
	pi.on("before_agent_start", async (event) => {
		if (!personalityContent) return undefined;

		if (isFirstTurn) {
			isFirstTurn = false;
			return {
				systemPrompt: event.systemPrompt + wrapPersonality(personalityContent),
			};
		}

		return {
			message: {
				customType: "system-reminder",
				content: wrapPersonality(personalityContent),
				display: false,
			},
		};
	});

	// 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;
				clearPersonality();
				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;
				}
				switchPersonality(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;
					clearPersonality();
					ctx.ui.notify(`Personality cleared (was: ${was})`, "info");
				}
				return;
			}

			switchPersonality(picked);
			ctx.ui.notify(`Set personality to "${picked}"`, "info");
		},
	});
}