import * as fs from "node:fs";
import * as path from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
	DEFAULT_CONFIG,
	enforceMutationSafeLimits,
	isRecord,
	loadAgentInventory,
	resolveAgentScope,
	type PolicyConfig,
	type SubagentInput,
} from "../../src/policy.ts";

const CONFIG_PATH = path.join(import.meta.dirname, "config.json");
const AGENT_DIR = process.env.PI_CODING_AGENT_DIR ?? path.join(process.env.HOME ?? "", ".pi", "agent");

function loadConfig(): PolicyConfig {
	try {
		const parsed = JSON.parse(fs.readFileSync(CONFIG_PATH, "utf8")) as Partial<PolicyConfig>;
		if (parsed.version !== 1) return DEFAULT_CONFIG;
		return {
			version: 1,
			enabled: parsed.enabled !== false,
			allowExplicitMutationTimeout: parsed.allowExplicitMutationTimeout === true,
			unknownAgentMayMutate: parsed.unknownAgentMayMutate !== false,
		};
	} catch {
		return DEFAULT_CONFIG;
	}
}

function isExecutionOrSchedule(input: SubagentInput): boolean {
	const action = typeof input.action === "string" ? input.action : undefined;
	return !action || action === "schedule";
}

function auditMessage(input: { removed: string[]; checkpointPromptsAdded: number; assessment: { reason: string } }, action: string | undefined): string {
	return `Mutation-safe subagent policy removed ${input.removed.join(", ")} for ${action === "schedule" ? "a scheduled" : "a"} mutation-capable run (${input.assessment.reason}). Added ${input.checkpointPromptsAdded} checkpoint instruction${input.checkpointPromptsAdded === 1 ? "" : "s"}. The run remains controllable with status, steer, interrupt, and stop; use subagent_wait timeout only to stop waiting without killing it.`;
}

export default function registerSubagentTimeoutSafety(pi: ExtensionAPI): void {
	// pi-subagents does not load its normal parent extension in ordinary child
	// processes. Avoid changing fanout-child semantics if this package is supplied
	// by an inherited extension configuration.
	if (process.env.PI_SUBAGENT_CHILD === "1") return;

	pi.on("tool_call", (event, ctx) => {
		if (event.toolName !== "subagent" || !isRecord(event.input)) return undefined;
		const input = event.input as SubagentInput;
		if (!isExecutionOrSchedule(input)) return undefined;

		const outcome = enforceMutationSafeLimits(input, loadAgentInventory(ctx.cwd, AGENT_DIR, resolveAgentScope(input.agentScope)), loadConfig());
		if (!outcome) return undefined;

		const action = typeof input.action === "string" ? input.action : undefined;
		const message = auditMessage(outcome, action);
		if (ctx.hasUI) ctx.ui.notify(message, "warning");
		pi.sendMessage({
			customType: "subagent-timeout-safety",
			content: message,
			display: true,
			details: {
				removed: outcome.removed,
				checkpointPromptsAdded: outcome.checkpointPromptsAdded,
				reason: outcome.assessment.reason,
				action: action ?? "execute",
			},
		});
		return undefined;
	});

	pi.registerCommand("subagent-timeout-safety", {
		description: "Show the mutation-safe subagent timeout policy",
		handler: async (_args, ctx) => {
			const policy = loadConfig();
			const inventory = loadAgentInventory(ctx.cwd, AGENT_DIR, "both");
			const unsafeDefaults = inventory.unsafeDefaults.length === 0
				? "none"
				: inventory.unsafeDefaults.map((entry) => `${entry.name} (${entry.limits.join(", ")}): ${entry.filePath}`).join("\n  - ");
			const lines = [
				"Mutation-safe subagent timeout policy",
				`- Enabled: ${policy.enabled}`,
				`- Explicit writer hard timeouts allowed: ${policy.allowExplicitMutationTimeout} (unsafe when true)`,
				`- Unknown agents fail closed as writers: ${policy.unknownAgentMayMutate}`,
				`- Agent definitions inspected: ${inventory.capabilities.size}`,
				"- Effect: removes caller-supplied run-level timeoutMs, maxRuntimeMs, and turnBudget only when the execution may mutate files.",
				"- Explicitly read-only tasks and known tool-restricted custom agents retain caller-supplied limits.",
				"- Applies to LLM-issued subagent tool calls, including action='schedule'. Direct pi-subagents slash/delegation/RPC/fanout executor paths remain outside Pi tool_call interception.",
				"- Manual stop/interrupt and protocol safety guards remain available.",
				`- Mutation-capable local-agent defaults that still need removal: ${unsafeDefaults}`,
				`- Config: ${CONFIG_PATH}`,
			].join("\n");
			ctx.ui.notify(lines, "info");
			pi.sendMessage({ customType: "subagent-timeout-safety", content: lines, display: true, details: policy });
		},
	});
}