import assert from "node:assert/strict";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import test from "node:test";

import {
	DEFAULT_CONFIG,
	enforceMutationSafeLimits,
	executionMayMutate,
	injectMutationSafetyHints,
	loadAgentInventory,
	removeHardLimits,
	type AgentInventory,
} from "../src/policy.ts";

function inventory(entries: Array<[string, boolean]>): AgentInventory {
	return {
		capabilities: new Map(entries.map(([name, mayMutate]) => [name, { name, mayMutate, known: true }])),
		unsafeDefaults: [],
	};
}

const strictPolicy = { ...DEFAULT_CONFIG };

test("removes every writer hard limit and adds a checkpoint", () => {
	const input = {
		agent: "writer",
		task: "Implement the migration.",
		timeoutMs: 900_000,
		maxRuntimeMs: 900_000,
		turnBudget: { maxTurns: 8 },
	};
	const result = enforceMutationSafeLimits(input, inventory([["writer", true]]), strictPolicy);
	assert.deepEqual(result?.removed.sort(), ["maxRuntimeMs", "timeoutMs", "turnBudget"]);
	assert.equal(result?.checkpointPromptsAdded, 1);
	assert.equal("timeoutMs" in input, false);
	assert.equal("maxRuntimeMs" in input, false);
	assert.equal("turnBudget" in input, false);
	assert.match(input.task, /Mutation-safe deadline policy is active/);
});

test("still removes limits for a writer that claims read-only intent", () => {
	const input = {
		agent: "writer",
		task: "Read-only review only. Do not modify files.",
		timeoutMs: 120_000,
	};
	const result = enforceMutationSafeLimits(input, inventory([["writer", true]]), strictPolicy);
	assert.deepEqual(result?.removed, ["timeoutMs"]);
	assert.match(result?.assessment.reason ?? "", /mutation-capable tools/);
});

test("known tool-restricted custom agents can retain limits", () => {
	const input = { agent: "reader", task: "Inspect the diff and report findings.", timeoutMs: 120_000 };
	assert.equal(executionMayMutate(input, inventory([["reader", false]]), strictPolicy).mayMutate, false);
	assert.equal(enforceMutationSafeLimits(input, inventory([["reader", false]]), strictPolicy), undefined);
});

test("unknown agents fail closed even when their task sounds observational", () => {
	const input = { agent: "not-installed", task: "Inspect the diff.", timeoutMs: 120_000 };
	const result = enforceMutationSafeLimits(input, inventory([]), strictPolicy);
	assert.deepEqual(result?.removed, ["timeoutMs"]);
	assert.match(result?.assessment.reason ?? "", /unknown and fails closed/);
});

test("mixed chain removes a shared hard limit and checkpoints concrete steps", () => {
	const input = {
		chain: [
			{ agent: "reader", task: "Read-only review only; no edits." },
			{ agent: "writer", task: "Apply the approved change." },
		],
		timeoutMs: 120_000,
	};
	const result = enforceMutationSafeLimits(input, inventory([["reader", false], ["writer", true]]), strictPolicy);
	assert.deepEqual(result?.removed, ["timeoutMs"]);
	assert.equal(result?.checkpointPromptsAdded, 2);
	assert.match(input.chain[0].task, /Mutation-safe deadline policy is active/);
	assert.match(input.chain[1].task, /Mutation-safe deadline policy is active/);
});

test("dynamic fanout fails closed and checkpoints its concrete template", () => {
	const input = {
		chain: [{
			expand: { from: { output: "targets", path: "/items" }, maxItems: 3 },
			parallel: { agent: "writer", task: "Apply item {item}." },
			collect: { as: "results" },
		}],
		timeoutMs: 120_000,
	};
	const result = enforceMutationSafeLimits(input, inventory([["writer", true]]), strictPolicy);
	assert.deepEqual(result?.removed, ["timeoutMs"]);
	assert.equal(result?.checkpointPromptsAdded, 1);
	assert.match(input.chain[0].parallel.task, /Mutation-safe deadline policy is active/);
	assert.equal("timeoutMs" in input, false);
});

test("policy helpers are idempotent", () => {
	const input = { agent: "writer", task: "Implement.", timeoutMs: 1 };
	assert.deepEqual(removeHardLimits(input), ["timeoutMs"]);
	assert.deepEqual(removeHardLimits(input), []);
	assert.equal(injectMutationSafetyHints(input), 1);
	assert.equal(injectMutationSafetyHints(input), 0);
});

test("agent inventory honors scope, recursively mirrors local discovery, and exposes unsafe defaults", () => {
	const root = fs.mkdtempSync(path.join(os.tmpdir(), "pi-timeout-policy-"));
	const priorExtraDirectories = process.env.PI_SUBAGENT_EXTRA_AGENT_DIRS;
	try {
		const agentDir = path.join(root, "agent-home");
		const cwd = path.join(root, "project", "nested");
		fs.mkdirSync(path.join(agentDir, "agents", "nested"), { recursive: true });
		const extraAgentDir = path.join(root, "extra-agents");
		fs.mkdirSync(extraAgentDir, { recursive: true });
		fs.mkdirSync(path.join(root, "project", ".pi", "agents"), { recursive: true });
		fs.mkdirSync(cwd, { recursive: true });
		fs.writeFileSync(path.join(agentDir, "agents", "nested", "reader.md"), "---\nname: reader\ndescription: reader\ntools: read, grep\n---\n");
		fs.writeFileSync(path.join(root, "project", ".pi", "agents", "writer.md"), "---\nname: writer\ndescription: writer\ntools: read, write\ntimeoutMs: 900000\nturnBudget: {\"maxTurns\": 8}\n---\n");
		fs.writeFileSync(path.join(extraAgentDir, "extra-reader.md"), "---\nname: extra-reader\ndescription: reader\ntools: read, grep\n---\n");
		process.env.PI_SUBAGENT_EXTRA_AGENT_DIRS = extraAgentDir;
		const actual = loadAgentInventory(cwd, agentDir);
		const userOnly = loadAgentInventory(cwd, agentDir, "user");
		const projectOnly = loadAgentInventory(cwd, agentDir, "project");
		assert.equal(actual.capabilities.get("reader")?.mayMutate, false);
		assert.equal(actual.capabilities.get("extra-reader")?.mayMutate, false);
		assert.equal(actual.capabilities.get("writer")?.mayMutate, true);
		assert.equal(userOnly.capabilities.has("writer"), false);
		assert.equal(projectOnly.capabilities.has("reader"), false);
		assert.deepEqual(actual.unsafeDefaults.map((entry) => ({ name: entry.name, limits: entry.limits })), [{ name: "writer", limits: ["timeoutMs", "turnBudget"] }]);
	} finally {
		if (priorExtraDirectories === undefined) delete process.env.PI_SUBAGENT_EXTRA_AGENT_DIRS;
		else process.env.PI_SUBAGENT_EXTRA_AGENT_DIRS = priorExtraDirectories;
		fs.rmSync(root, { recursive: true, force: true });
	}
});