/**
 * Diagnostic budget guard for unusually long auto-mode units.
 *
 * This is intentionally not a blind tool-count kill switch. It gives the agent
 * explicit turns to explain whether the unit is legitimately large, stuck, or
 * churning, then pauses only if the unit keeps consuming budget afterward.
 */

import { execFileSync } from "node:child_process";
import { createHash } from "node:crypto";
import { existsSync, lstatSync, readdirSync, readFileSync } from "node:fs";
import type { ExtensionContext } from "@singularity-forge/pi-coding-agent";
import { formatTokenCount } from "../shared/format-utils.js";
import type { AutoSupervisorConfig } from "./preferences-types.js";

export const DEFAULT_RUNAWAY_TOOL_CALL_WARNING = 60;
export const DEFAULT_RUNAWAY_TOKEN_WARNING = 1_000_000;
export const DEFAULT_RUNAWAY_ELAPSED_MINUTES = 20;
export const DEFAULT_RUNAWAY_CHANGED_FILES_WARNING = 75;
export const DEFAULT_RUNAWAY_DIAGNOSTIC_TURNS = 2;
export const DEFAULT_RUNAWAY_MIN_INTERVAL_MS = 120_000;
const EXECUTE_NO_PROGRESS_TOOL_WARNING = 25;
const EXECUTE_NO_PROGRESS_TOKEN_WARNING = 500_000;
const DURABLE_SF_ARTIFACT_PATHS = [
	".sf/milestones",
	".sf/approvals",
	".sf/DECISIONS.md",
	".sf/KNOWLEDGE.md",
	".sf/STATE.md",
];

export interface RunawayGuardConfig {
	enabled: boolean;
	toolCallWarning: number;
	tokenWarning: number;
	elapsedMs: number;
	changedFilesWarning: number;
	diagnosticTurns: number;
	hardPause: boolean;
	minIntervalMs: number;
}

export interface RunawayGuardMetrics {
	toolCalls: number;
	sessionTokens: number;
	elapsedMs: number;
	changedFiles?: number;
	worktreeFingerprint?: string | null;
	worktreeChangedSinceStart?: boolean;
	topTools?: Record<string, number>;
}

export interface RunawayGuardPauseMetadata {
	reason: string;
	pausedAt: number;
	unitType: string;
	unitId: string;
	diagnosticTurns: number;
	warningsSent: number;
	thresholdReasons: string[];
	metrics: RunawayGuardMetrics;
	lastWarningMetrics: Pick<
		RunawayGuardMetrics,
		"toolCalls" | "sessionTokens" | "elapsedMs"
	>;
	thresholds: Pick<
		RunawayGuardConfig,
		| "toolCallWarning"
		| "tokenWarning"
		| "elapsedMs"
		| "changedFilesWarning"
		| "minIntervalMs"
	>;
}

interface RunawayGuardState {
	unitKey: string;
	baselineSessionTokens: number;
	baselineChangedFiles: number;
	baselineWorktreeFingerprint: string | null;
	warningsSent: number;
	lastWarningAt: number;
	lastToolCalls: number;
	lastSessionTokens: number;
	lastElapsedMs: number;
	finalWarningSent: boolean;
}

let state: RunawayGuardState | null = null;

export function resetRunawayGuardState(
	unitType: string,
	unitId: string,
	baseline?: {
		sessionTokens?: number;
		changedFiles?: number;
		worktreeFingerprint?: string | null;
	},
): void {
	state = {
		unitKey: `${unitType}/${unitId}`,
		baselineSessionTokens: baseline?.sessionTokens ?? 0,
		baselineChangedFiles: baseline?.changedFiles ?? 0,
		baselineWorktreeFingerprint: baseline?.worktreeFingerprint ?? null,
		warningsSent: 0,
		lastWarningAt: 0,
		lastToolCalls: 0,
		lastSessionTokens: 0,
		lastElapsedMs: 0,
		finalWarningSent: false,
	};
}

export function clearRunawayGuardState(): void {
	state = null;
}

export function resolveRunawayGuardConfig(
	supervisor: AutoSupervisorConfig | undefined,
): RunawayGuardConfig {
	return {
		enabled: supervisor?.runaway_guard_enabled !== false,
		toolCallWarning:
			supervisor?.runaway_tool_call_warning ??
			DEFAULT_RUNAWAY_TOOL_CALL_WARNING,
		tokenWarning:
			supervisor?.runaway_token_warning ?? DEFAULT_RUNAWAY_TOKEN_WARNING,
		elapsedMs:
			(supervisor?.runaway_elapsed_minutes ?? DEFAULT_RUNAWAY_ELAPSED_MINUTES) *
			60 *
			1000,
		changedFilesWarning:
			supervisor?.runaway_changed_files_warning ??
			DEFAULT_RUNAWAY_CHANGED_FILES_WARNING,
		diagnosticTurns:
			supervisor?.runaway_diagnostic_turns ?? DEFAULT_RUNAWAY_DIAGNOSTIC_TURNS,
		hardPause: supervisor?.runaway_hard_pause !== false,
		minIntervalMs: DEFAULT_RUNAWAY_MIN_INTERVAL_MS,
	};
}

export function collectSessionTokenUsage(ctx: ExtensionContext): number {
	try {
		const entries = ctx.sessionManager?.getEntries?.() ?? [];
		let total = 0;
		for (const entry of entries as unknown as Array<{
			message?: { role?: string; usage?: Record<string, unknown> };
		}>) {
			const message = entry.message;
			if (message?.role !== "assistant" || !message.usage) continue;
			const usage = message.usage;
			const totalTokens = numeric(usage.totalTokens ?? usage.total);
			if (totalTokens > 0) {
				total += totalTokens;
				continue;
			}
			total +=
				numeric(usage.input) +
				numeric(usage.output) +
				numeric(usage.cacheRead) +
				numeric(usage.cacheWrite);
		}
		return total;
	} catch {
		return 0;
	}
}

export function countChangedFiles(cwd: string): number {
	try {
		const out = execFileSync("git", ["status", "--short"], {
			cwd,
			encoding: "utf8",
			stdio: ["ignore", "pipe", "ignore"],
			timeout: 2000,
		});
		return out
			.split("\n")
			.map((line) => line.trim())
			.filter(Boolean).length;
	} catch {
		return 0;
	}
}

export function collectWorktreeFingerprint(cwd: string): string | null {
	try {
		const status = execFileSync(
			"git",
			["status", "--porcelain=v1", "--untracked-files=all"],
			{
				cwd,
				encoding: "utf8",
				stdio: ["ignore", "pipe", "ignore"],
				timeout: 2000,
			},
		);
		const lines = status
			.split("\n")
			.map((line) => line.trimEnd())
			.filter(Boolean);

		const hash = createHash("sha256");
		if (lines.length === 0) {
			hash.update("git-clean");
			hash.update("\0");
		}
		for (const line of lines) {
			hash.update(line);
			hash.update("\0");
			const filePath = parsePorcelainPath(line);
			if (!filePath) continue;
			appendFileFingerprint(hash, cwd, filePath);
		}
		appendDurableSfArtifactFingerprint(hash, cwd);
		return hash.digest("hex");
	} catch {
		return null;
	}
}

function appendDurableSfArtifactFingerprint(
	hash: ReturnType<typeof createHash>,
	cwd: string,
): void {
	hash.update("sf-artifacts");
	hash.update("\0");
	for (const artifactPath of DURABLE_SF_ARTIFACT_PATHS) {
		appendPathFingerprint(hash, cwd, artifactPath);
	}
}

function appendPathFingerprint(
	hash: ReturnType<typeof createHash>,
	cwd: string,
	relativePath: string,
): void {
	const fullPath = `${cwd}/${relativePath}`;
	if (!existsSync(fullPath)) {
		hash.update(`missing:${relativePath}`);
		hash.update("\0");
		return;
	}

	let stat: ReturnType<typeof lstatSync>;
	try {
		stat = lstatSync(fullPath);
	} catch {
		hash.update(`unreadable:${relativePath}`);
		hash.update("\0");
		return;
	}

	if (stat.isDirectory()) {
		hash.update(`dir:${relativePath}`);
		hash.update("\0");
		let entries: string[];
		try {
			entries = readdirSync(fullPath).sort();
		} catch {
			hash.update(`unreadable-dir:${relativePath}`);
			hash.update("\0");
			return;
		}
		for (const entry of entries) {
			appendPathFingerprint(hash, cwd, `${relativePath}/${entry}`);
		}
		return;
	}

	appendFileFingerprint(hash, cwd, relativePath);
}

function appendFileFingerprint(
	hash: ReturnType<typeof createHash>,
	cwd: string,
	relativePath: string,
): void {
	try {
		const stat = lstatSync(`${cwd}/${relativePath}`);
		if (!stat.isFile()) {
			hash.update(
				`type:${relativePath}:${stat.isDirectory() ? "dir" : "other"}`,
			);
			hash.update("\0");
			return;
		}
		hash.update(`file:${relativePath}`);
		hash.update("\0");
		hash.update(readFileSync(`${cwd}/${relativePath}`));
		hash.update("\0");
	} catch {
		hash.update(`unreadable-or-deleted:${relativePath}`);
		hash.update("\0");
	}
}

export type RunawayGuardDecision =
	| { action: "none" }
	| { action: "warn"; message: string; final: boolean }
	| { action: "pause"; reason: string; metadata: RunawayGuardPauseMetadata };

export function evaluateRunawayGuard(
	unitType: string,
	unitId: string,
	metrics: RunawayGuardMetrics,
	config: RunawayGuardConfig,
	now = Date.now(),
): RunawayGuardDecision {
	if (!config.enabled) return { action: "none" };
	if (config.diagnosticTurns <= 0) return { action: "none" };

	const unitKey = `${unitType}/${unitId}`;
	if (!state || state.unitKey !== unitKey)
		resetRunawayGuardState(unitType, unitId);
	const s = state!;
	const unitMetrics = normalizeMetricsToUnit(metrics, s);

	const reasons = thresholdReasons(unitType, unitMetrics, config);
	if (reasons.length === 0) return { action: "none" };

	if (
		s.lastWarningAt > 0 &&
		now - s.lastWarningAt < config.minIntervalMs &&
		!hasMeaningfulGrowth(unitMetrics, s, config)
	) {
		return { action: "none" };
	}

	if (
		config.hardPause &&
		s.finalWarningSent &&
		hasMeaningfulGrowth(unitMetrics, s, config)
	) {
		const reason =
			`Runaway guard paused ${unitType} ${unitId}: budget kept growing after ` +
			`${config.diagnosticTurns} diagnostic turn(s). ` +
			formatMetricSummary(unitMetrics);
		return {
			action: "pause",
			reason,
			metadata: {
				reason,
				pausedAt: now,
				unitType,
				unitId,
				diagnosticTurns: config.diagnosticTurns,
				warningsSent: s.warningsSent,
				thresholdReasons: reasons,
				metrics: unitMetrics,
				lastWarningMetrics: {
					toolCalls: s.lastToolCalls,
					sessionTokens: s.lastSessionTokens,
					elapsedMs: s.lastElapsedMs,
				},
				thresholds: {
					toolCallWarning: config.toolCallWarning,
					tokenWarning: config.tokenWarning,
					elapsedMs: config.elapsedMs,
					changedFilesWarning: config.changedFilesWarning,
					minIntervalMs: config.minIntervalMs,
				},
			},
		};
	}

	const final = s.warningsSent + 1 >= config.diagnosticTurns;
	s.warningsSent++;
	s.lastWarningAt = now;
	s.lastToolCalls = unitMetrics.toolCalls;
	s.lastSessionTokens = unitMetrics.sessionTokens;
	s.lastElapsedMs = unitMetrics.elapsedMs;
	if (final) s.finalWarningSent = true;

	return {
		action: "warn",
		final,
		message: buildRunawayGuardMessage(
			unitType,
			unitId,
			unitMetrics,
			reasons,
			final,
		),
	};
}

function normalizeMetricsToUnit(
	metrics: RunawayGuardMetrics,
	state: RunawayGuardState,
): RunawayGuardMetrics {
	const worktreeChangedSinceStart =
		metrics.worktreeFingerprint !== undefined &&
		metrics.worktreeFingerprint !== null &&
		state.baselineWorktreeFingerprint !== null
			? metrics.worktreeFingerprint !== state.baselineWorktreeFingerprint
			: metrics.worktreeChangedSinceStart;
	return {
		...metrics,
		sessionTokens: Math.max(
			0,
			metrics.sessionTokens - state.baselineSessionTokens,
		),
		changedFiles:
			metrics.changedFiles === undefined
				? undefined
				: Math.max(0, metrics.changedFiles - state.baselineChangedFiles),
		worktreeChangedSinceStart,
	};
}

function thresholdReasons(
	unitType: string,
	metrics: RunawayGuardMetrics,
	config: RunawayGuardConfig,
): string[] {
	const reasons: string[] = [];
	if (
		config.toolCallWarning > 0 &&
		metrics.toolCalls >= config.toolCallWarning
	) {
		reasons.push(
			`${metrics.toolCalls} tool calls (warning ${config.toolCallWarning})`,
		);
	}
	if (config.tokenWarning > 0 && metrics.sessionTokens >= config.tokenWarning) {
		reasons.push(
			`${formatTokenCount(metrics.sessionTokens)} session tokens (warning ${formatTokenCount(config.tokenWarning)})`,
		);
	}
	if (config.elapsedMs > 0 && metrics.elapsedMs >= config.elapsedMs) {
		reasons.push(
			`${Math.round(metrics.elapsedMs / 60000)}min elapsed (warning ${Math.round(config.elapsedMs / 60000)}min)`,
		);
	}
	if (
		config.changedFilesWarning > 0 &&
		(metrics.changedFiles ?? 0) >= config.changedFilesWarning
	) {
		reasons.push(
			`${metrics.changedFiles} new changed files (warning ${config.changedFilesWarning})`,
		);
	}
	if (
		unitType === "execute-task" &&
		(metrics.changedFiles ?? 0) === 0 &&
		metrics.worktreeChangedSinceStart !== true &&
		metrics.toolCalls >= EXECUTE_NO_PROGRESS_TOOL_WARNING &&
		metrics.sessionTokens >= EXECUTE_NO_PROGRESS_TOKEN_WARNING
	) {
		reasons.push(
			`no new file changes after ${metrics.toolCalls} tool calls and ${formatTokenCount(metrics.sessionTokens)} tokens`,
		);
	}
	return reasons;
}

function hasMeaningfulGrowth(
	metrics: RunawayGuardMetrics,
	state: RunawayGuardState,
	config: RunawayGuardConfig,
): boolean {
	const toolGrowth = Math.max(5, Math.floor(config.toolCallWarning / 4));
	const tokenGrowth = Math.max(50_000, Math.floor(config.tokenWarning / 4));
	return (
		metrics.toolCalls - state.lastToolCalls >= toolGrowth ||
		metrics.sessionTokens - state.lastSessionTokens >= tokenGrowth ||
		metrics.elapsedMs - state.lastElapsedMs >= config.minIntervalMs
	);
}

function buildRunawayGuardMessage(
	unitType: string,
	unitId: string,
	metrics: RunawayGuardMetrics,
	reasons: string[],
	final: boolean,
): string {
	const topTools = metrics.topTools
		? Object.entries(metrics.topTools)
				.sort(([, a], [, b]) => b - a)
				.slice(0, 5)
				.map(([name, count]) => `${name}x${count}`)
				.join(", ")
		: "";
	const title = final
		? "**RUNAWAY UNIT FINAL WARNING - write diagnosis and handoff now.**"
		: "**RUNAWAY UNIT BUDGET WARNING - diagnose before continuing.**";
	return [
		title,
		`Unit: ${unitType} ${unitId}`,
		`Budget signals: ${reasons.join("; ")}.`,
		topTools ? `Tool mix: ${topTools}.` : "",
		formatChangedFilesLine(unitType, metrics),
		"",
		final
			? "You have already received a budget warning. Do not start new exploration. Write or update the durable artifact/handoff now, explicitly stating whether the unit was legitimately large, blocked, or stuck in a loop."
			: "Before more exploration or broad edits, state why this unit is still running: legitimately large, blocked, or stuck/churning. Then either finish the required artifact or write a precise handoff.",
	]
		.filter(Boolean)
		.join("\n");
}

function formatChangedFilesLine(
	unitType: string,
	metrics: RunawayGuardMetrics,
): string {
	if ((metrics.changedFiles ?? 0) > 0) {
		return `Working tree has ${metrics.changedFiles} new changed file(s) since this unit started. Active edits are not automatically healthy progress; check for repeated or broad churn.`;
	}
	if (unitType === "execute-task" && metrics.worktreeChangedSinceStart) {
		return "Working tree has 0 new changed file paths, but dirty file content changed since this execute-task started.";
	}
	if (unitType === "execute-task") {
		return "Working tree has 0 new changed files since this execute-task started. For implementation work, that is no durable progress yet.";
	}
	return "";
}

function formatMetricSummary(metrics: RunawayGuardMetrics): string {
	return [
		`${metrics.toolCalls} tool calls`,
		`${formatTokenCount(metrics.sessionTokens)} tokens`,
		`${Math.round(metrics.elapsedMs / 60000)}min elapsed`,
		metrics.changedFiles !== undefined
			? `${metrics.changedFiles} new changed files`
			: "",
		metrics.worktreeChangedSinceStart ? "dirty file content changed" : "",
	]
		.filter(Boolean)
		.join(", ");
}

function parsePorcelainPath(line: string): string | null {
	if (line.length < 4) return null;
	let filePath = line.slice(3);
	const renameSeparator = " -> ";
	if (filePath.includes(renameSeparator)) {
		filePath = filePath.slice(
			filePath.lastIndexOf(renameSeparator) + renameSeparator.length,
		);
	}
	if (filePath.startsWith('"') && filePath.endsWith('"')) {
		filePath = filePath.slice(1, -1);
	}
	return filePath || null;
}

function numeric(value: unknown): number {
	return typeof value === "number" && Number.isFinite(value) ? value : 0;
}
