import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";

export interface PolicyConfig {
	version: 1;
	enabled: boolean;
	/** Explicitly unsafe escape hatch. Keep false in normal operation. */
	allowExplicitMutationTimeout: boolean;
	/** Treat unresolved agent names as writers. */
	unknownAgentMayMutate: boolean;
}

export interface AgentCapabilities {
	name: string;
	mayMutate: boolean;
	known: boolean;
}

export interface AgentInventory {
	capabilities: Map<string, AgentCapabilities>;
	unsafeDefaults: Array<{ name: string; filePath: string; limits: string[] }>;
}

export interface MutationAssessment {
	mayMutate: boolean;
	reason: string;
}

export type AgentScope = "user" | "project" | "both";
export type SubagentInput = Record<string, unknown>;
export type ExecutionItem = Record<string, unknown>;

export const DEFAULT_CONFIG: PolicyConfig = {
	version: 1,
	enabled: true,
	allowExplicitMutationTimeout: false,
	unknownAgentMayMutate: true,
};

const READ_ONLY_TOOL_NAMES = new Set([
	"read",
	"grep",
	"find",
	"ls",
	"web_search",
	"fetch_content",
	"get_search_content",
	"intercom",
	"contact_supervisor",
	"structured_output",
	"subagent_wait",
]);

const READ_ONLY_TASK_PATTERNS = [
	/\bread[ -]?only\b/i,
	/\bdo not modify\b/i,
	/\bdo not edit\b/i,
	/\bno edits?\b/i,
	/\bdo not write\b/i,
	/\bwithout (?:making )?(?:any )?(?:edits?|changes?|modifications?)\b/i,
	/\b(?:inspect|review)\s+only\b/i,
	/\bonly\s+(?:inspect|review)\b/i,
	/\b不要修改\b/u,
	/\b不修改\b/u,
	/\b禁止修改\b/u,
	/\b仅(?:做)?(?:审查|检查|分析)\b/u,
	/\b只读\b/u,
];

export function isRecord(value: unknown): value is Record<string, unknown> {
	return Boolean(value) && typeof value === "object" && !Array.isArray(value);
}

export function parseTools(raw: string | undefined): string[] {
	if (!raw) return [];
	return raw.split(",").map((tool) => tool.trim()).filter(Boolean);
}

export function parseFrontmatter(source: string): Record<string, string> {
	const match = source.match(/^---\r?\n([\s\S]*?)\r?\n---(?:\r?\n|$)/);
	if (!match) return {};
	const fields: Record<string, string> = {};
	for (const line of match[1].split(/\r?\n/)) {
		const separator = line.indexOf(":");
		if (separator < 1) continue;
		const key = line.slice(0, separator).trim();
		const value = line.slice(separator + 1).trim();
		if (key) fields[key] = value;
	}
	return fields;
}

function agentDirectories(cwd: string, agentDir: string, scope: AgentScope): string[] {
	const projectRoot = findNearestProjectRoot(cwd);
	const userDirs = scope === "project"
		? []
		: [...extraUserAgentDirectories(), path.join(agentDir, "agents"), path.join(os.homedir(), ".agents")];
	const projectDirs = scope === "user" || !projectRoot
		? []
		: [path.join(projectRoot, ".agents"), path.join(projectRoot, ".pi", "agents")];
	// Later paths are higher precedence: old user < new user < legacy project < .pi project.
	return [...userDirs, ...projectDirs];
}

function extraUserAgentDirectories(): string[] {
	const raw = process.env.PI_SUBAGENT_EXTRA_AGENT_DIRS;
	if (!raw) return [];
	return raw.split(path.delimiter).map((directory) => directory.trim()).filter(Boolean);
}

function isDirectory(candidate: string): boolean {
	try {
		return fs.statSync(candidate).isDirectory();
	} catch {
		return false;
	}
}

function findNearestProjectRoot(cwd: string): string | null {
	let current = path.resolve(cwd);
	while (true) {
		if (isDirectory(path.join(current, ".pi")) || isDirectory(path.join(current, ".agents"))) return current;
		const parent = path.dirname(current);
		if (parent === current) return null;
		current = parent;
	}
}

function toolsMayMutate(tools: string[]): boolean {
	return tools.length === 0 || tools.some((tool) => !READ_ONLY_TOOL_NAMES.has(tool));
}

function listMarkdownFilesRecursive(directory: string): string[] {
	const files: string[] = [];
	try {
		for (const entry of fs.readdirSync(directory, { withFileTypes: true })) {
			const entryPath = path.join(directory, entry.name);
			if (entry.isDirectory()) {
				files.push(...listMarkdownFilesRecursive(entryPath));
			} else if (entry.isFile() && entry.name.endsWith(".md") && !entry.name.endsWith(".chain.md")) {
				files.push(entryPath);
			}
		}
	} catch {
		// Missing/unreadable locations are handled safely by unknown-agent fallback.
	}
	return files.sort((left, right) => left.localeCompare(right));
}

/**
 * Mirrors pi-subagents' relevant discovery directories and precedence for capability
 * classification. It deliberately does not resolve package-agent sources: package
 * agents are unresolved here and therefore fail closed.
 */
export function loadAgentInventory(cwd: string, agentDir: string, scope: AgentScope = "both"): AgentInventory {
	const definitions = new Map<string, { capabilities: AgentCapabilities; filePath: string; limits: string[] }>();
	for (const directory of agentDirectories(cwd, agentDir, scope)) {
		for (const filePath of listMarkdownFilesRecursive(directory)) {
			try {
				const frontmatter = parseFrontmatter(fs.readFileSync(filePath, "utf8"));
				const name = frontmatter.name?.trim();
				if (!name) continue;
				const mayMutate = toolsMayMutate(parseTools(frontmatter.tools));
				const limits = ["timeoutMs", "turnBudget"].filter((field) => Boolean(frontmatter[field]?.trim()));
				definitions.set(name, { capabilities: { name, mayMutate, known: true }, filePath, limits });
			} catch {
				// A single unreadable definition does not prevent conservative fallback.
			}
		}
	}
	return {
		capabilities: new Map([...definitions].map(([name, definition]) => [name, definition.capabilities])),
		unsafeDefaults: [...definitions.values()]
			.filter((definition) => definition.capabilities.mayMutate && definition.limits.length > 0)
			.map(({ capabilities, filePath, limits }) => ({ name: capabilities.name, filePath, limits })),
	};
}

const BUILTIN_MUTATION_CAPABLE_ROLES = new Set(["worker", "reviewer", "researcher", "scout", "context-builder", "planner", "oracle", "delegate"]);

function builtinCapabilities(agent: string): AgentCapabilities | undefined {
	// In pi-subagents 0.35.1 every built-in role includes bash, write, or edit.
	return BUILTIN_MUTATION_CAPABLE_ROLES.has(agent) ? { name: agent, mayMutate: true, known: true } : undefined;
}

function resolveCapabilities(agent: string, inventory: AgentInventory, policy: PolicyConfig): AgentCapabilities {
	return inventory.capabilities.get(agent)
		?? builtinCapabilities(agent)
		?? { name: agent, mayMutate: policy.unknownAgentMayMutate, known: false };
}

export function itemMayMutate(item: ExecutionItem, inventory: AgentInventory, policy: PolicyConfig): MutationAssessment {
	const agent = typeof item.agent === "string" ? item.agent : undefined;
	if (!agent) return { mayMutate: true, reason: "execution item has no concrete agent" };
	const capabilities = resolveCapabilities(agent, inventory, policy);
	if (!capabilities.known && policy.unknownAgentMayMutate) return { mayMutate: true, reason: `${agent} is unknown and fails closed` };
	// Prompt text and acceptanceRole do not remove tools. A writer can still
	// mutate despite being asked to review, so only a tool-restricted agent may
	// retain a hard limit. This is the safety boundary, not task intent.
	if (capabilities.mayMutate) return { mayMutate: true, reason: `${agent} has mutation-capable tools` };
	return { mayMutate: false, reason: `${agent} is configured read-only` };
}

export function resolveAgentScope(value: unknown): AgentScope {
	return value === "user" || value === "project" ? value : "both";
}

export function executionMayMutate(input: SubagentInput, inventory: AgentInventory, policy: PolicyConfig): MutationAssessment {
	if (Array.isArray(input.tasks)) {
		for (const task of input.tasks) {
			if (!isRecord(task)) return { mayMutate: true, reason: "unrecognized parallel task" };
			const assessment = itemMayMutate(task, inventory, policy);
			if (assessment.mayMutate) return assessment;
		}
		return { mayMutate: false, reason: "every parallel task is read-only" };
	}
	if (Array.isArray(input.chain)) {
		for (const step of input.chain) {
			if (!isRecord(step)) return { mayMutate: true, reason: "unrecognized chain step" };
			if (Array.isArray(step.parallel)) {
				for (const task of step.parallel) {
					if (!isRecord(task)) return { mayMutate: true, reason: "unrecognized parallel task" };
					const assessment = itemMayMutate(task, inventory, policy);
					if (assessment.mayMutate) return assessment;
				}
				continue;
			}
			if (isRecord(step.parallel)) return { mayMutate: true, reason: "dynamic fanout can materialize mutation-capable work" };
			const assessment = itemMayMutate(step, inventory, policy);
			if (assessment.mayMutate) return assessment;
		}
		return { mayMutate: false, reason: "every chain step is read-only" };
	}
	if (typeof input.agent === "string") return itemMayMutate(input, inventory, policy);
	return { mayMutate: false, reason: "not an execution request" };
}

export function hasHardLimit(input: SubagentInput): boolean {
	return input.timeoutMs !== undefined || input.maxRuntimeMs !== undefined || input.turnBudget !== undefined;
}

export function removeHardLimits(input: SubagentInput): string[] {
	const removed: string[] = [];
	for (const key of ["timeoutMs", "maxRuntimeMs", "turnBudget"] as const) {
		if (input[key] !== undefined) {
			delete input[key];
			removed.push(key);
		}
	}
	return removed;
}

const CHECKPOINT_HINT = [
	"Mutation-safe deadline policy is active: this task has no hard wall-clock timeout or hard turn budget.",
	"Before a natural pause or after substantial work, report a checkpoint with changed files, validation state, and remaining work.",
	"If blocked, ask the supervisor or return a clear checkpoint instead of continuing indefinitely.",
].join(" ");

function injectCheckpoint(item: ExecutionItem): boolean {
	const task = typeof item.task === "string" ? item.task.trim() : "";
	if (!task || task.includes("Mutation-safe deadline policy is active")) return false;
	item.task = `${task}\n\n${CHECKPOINT_HINT}`;
	return true;
}

export function injectMutationSafetyHints(input: SubagentInput): number {
	let changed = 0;
	if (Array.isArray(input.tasks)) {
		for (const task of input.tasks) if (isRecord(task) && injectCheckpoint(task)) changed++;
		return changed;
	}
	if (Array.isArray(input.chain)) {
		for (const step of input.chain) {
			if (!isRecord(step)) continue;
			if (Array.isArray(step.parallel)) {
				for (const task of step.parallel) if (isRecord(task) && injectCheckpoint(task)) changed++;
			} else if (isRecord(step.parallel) && injectCheckpoint(step.parallel)) {
				changed++;
			} else if (typeof step.agent === "string" && injectCheckpoint(step)) {
				changed++;
			}
		}
		return changed;
	}
	return injectCheckpoint(input) ? 1 : 0;
}

export function enforceMutationSafeLimits(input: SubagentInput, inventory: AgentInventory, policy: PolicyConfig): { removed: string[]; checkpointPromptsAdded: number; assessment: MutationAssessment } | undefined {
	if (!hasHardLimit(input) || !policy.enabled || policy.allowExplicitMutationTimeout) return undefined;
	const assessment = executionMayMutate(input, inventory, policy);
	if (!assessment.mayMutate) return undefined;
	const removed = removeHardLimits(input);
	return { removed, checkpointPromptsAdded: injectMutationSafetyHints(input), assessment };
}