/**
 * Subagent pane lifecycle reminder.
 *
 * When the parent agent dispatches to a `pane: true` agent via `subagent`,
 * this extension floats a cleanup reminder to the parent only:
 *
 *  - immediately: appended to the `subagent` tool result of that call;
 *  - persistently: a one-line system-prompt note on later turns until the
 *    parent calls `stop_subagent` for that agent (or the session ends).
 *
 * Nothing is injected for bg agents (`pane: false`), for `delegate_subagent`
 * (pane targets are rejected there anyway), for `/agents` commands, or in
 * sessions that never spawn a pane agent. Other sessions and workflows stay
 * untouched.
 */

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

interface SubagentInput {
	agent?: unknown;
	tasks?: Array<{ agent?: unknown }>;
	chain?: Array<{ agent?: unknown }>;
}

const USER_AGENT_DIRS: string[] = [
	path.join(homedir(), ".pi", "agent", "agents"),
	path.join(homedir(), ".claude", "agents"),
];

function projectAgentDirs(cwd: string): string[] {
	const dirs: string[] = [];
	let current = path.resolve(cwd);
	for (;;) {
		dirs.push(path.join(current, ".pi", "agents"), path.join(current, ".claude", "agents"));
		const parent = path.dirname(current);
		if (parent === current) break;
		current = parent;
	}
	return dirs;
}

function hasPaneFrontmatter(text: string): boolean {
	const block = text.match(/^---\r?\n([\s\S]*?)\r?\n---/);
	return block ? /^\s*pane\s*:\s*(?:true|yes|on)\s*$/im.test(block[1]) : false;
}

function isPaneAgent(name: string, cwd: string): boolean {
	for (const dir of [...USER_AGENT_DIRS, ...projectAgentDirs(cwd)]) {
		for (const ext of [".md", ".markdown"]) {
			try {
				if (hasPaneFrontmatter(fs.readFileSync(path.join(dir, `${name}${ext}`), "utf8"))) return true;
			} catch {
				// not present in this directory
			}
		}
	}
	return false;
}

function agentNamesFrom(input: SubagentInput): string[] {
	const names = new Set<string>();
	const add = (value: unknown) => {
		if (typeof value === "string" && value.trim()) names.add(value.trim());
	};
	add(input.agent);
	for (const item of input.tasks ?? []) add(item.agent);
	for (const item of input.chain ?? []) add(item.agent);
	return [...names];
}

function reminderText(names: string[]): string {
	const quoted = names.map((name) => `\`${name}\``).join(", ");
	return (
		`\n\nPane cleanup: you spawned persistent pane agent(s) ${quoted}. ` +
		`After you collect the result, call \`stop_subagent\` with the agent name to kill its tmux pane. ` +
		`Skip only if you plan a follow-up (steer or another subagent call to the same agent); stop it at the end of the workflow. ` +
		`Stopping preserves the session file, so a later \`subagent\` call resumes the agent's memory.`
	);
}

export default function (pi: ExtensionAPI) {
	// Agent names whose panes are still alive this session.
	const pending = new Set<string>();

	pi.on("tool_call", (event) => {
		if (event.toolName === "stop_subagent") {
			const agent = (event.input as { agent?: unknown } | undefined)?.agent;
			if (typeof agent === "string") pending.delete(agent);
		}
	});

	pi.on("tool_result", (event, ctx) => {
		if (event.toolName !== "subagent" || event.isError) return;
		const input = (event.input as SubagentInput | undefined) ?? {};
		const paneAgents = agentNamesFrom(input).filter((name) => isPaneAgent(name, ctx.cwd));
		if (paneAgents.length === 0) return;
		for (const name of paneAgents) pending.add(name);
		return {
			content: [...(event.content ?? []), { type: "text" as const, text: reminderText(paneAgents) }],
		};
	});

	pi.on("before_agent_start", (event) => {
		if (pending.size === 0) return;
		const note =
			"\n\nPane cleanup (pending): " +
			`[${[...pending].join(", ")}]. ` +
			"After collecting their results, call stop_subagent for each (unless a follow-up is planned).";
		return { systemPrompt: event.systemPrompt + note };
	});
}