/**
 * pi-continual: a minimal self-improving harness for pi.
 *
 * - `repl` tool: persistent Python kernel with rlm() sub-agents, harness CRUD, history()
 * - harness state (.pi/harness/{memory,prompts}) is injected into the system prompt each turn
 * - `/refine` (and a `refine` tool): background self-refinement of the harness state
 * - `/goal` + `/gate` + `goal_complete` tool: bounded autonomous mode
 */
import { type ChildProcess, spawn } from "node:child_process";
import { existsSync, mkdirSync, readFileSync, readdirSync, writeFileSync } from "node:fs";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import { Type } from "typebox";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

const EXT_DIR = dirname(fileURLToPath(import.meta.url));
const KERNEL_PY = join(EXT_DIR, "kernel.py");
const MARKER = "\x00PIC\x00";

const MAX_OUTPUT = 30_000;

function truncate(text: string, max = MAX_OUTPUT): string {
	if (text.length <= max) return text;
	return `${text.slice(0, max / 2)}\n...[truncated ${text.length - max} chars]...\n${text.slice(-max / 2)}`;
}

class Kernel {
	private proc: ChildProcess;
	private buffer = "";
	private pending = new Map<string, { resolve: (r: { ok: boolean; output: string }) => void; reject: (e: Error) => void }>();
	private nextId = 0;

	constructor(cwd: string, sessionFile: string | undefined) {
		this.proc = spawn("python3", [KERNEL_PY], {
			cwd,
			stdio: ["pipe", "pipe", "pipe"],
			env: {
				...process.env,
				PI_CONTINUAL_ROOT: cwd,
				PI_CONTINUAL_SESSION: sessionFile ?? "",
			},
		});
		this.proc.stdout?.on("data", (chunk: Buffer) => {
			this.buffer += chunk.toString("utf8");
			let idx: number;
			while ((idx = this.buffer.indexOf("\n")) >= 0) {
				const line = this.buffer.slice(0, idx);
				this.buffer = this.buffer.slice(idx + 1);
				if (!line.startsWith(MARKER)) continue;
				try {
					const msg = JSON.parse(line.slice(MARKER.length));
					const waiter = this.pending.get(msg.id);
					if (waiter) {
						this.pending.delete(msg.id);
						waiter.resolve({ ok: msg.ok, output: msg.output ?? "" });
					}
				} catch {
					// ignore malformed protocol lines
				}
			}
		});
		this.proc.on("exit", () => {
			for (const waiter of this.pending.values()) {
				waiter.reject(new Error("REPL kernel exited (state lost; next call starts a fresh kernel)"));
			}
			this.pending.clear();
		});
	}

	get alive(): boolean {
		return this.proc.exitCode === null && !this.proc.killed;
	}

	execute(code: string, timeoutMs: number, signal?: AbortSignal): Promise<{ ok: boolean; output: string }> {
		const id = String(this.nextId++);
		return new Promise((resolve, reject) => {
			const timer = setTimeout(() => {
				this.pending.delete(id);
				this.dispose();
				reject(new Error(`REPL execution timed out after ${timeoutMs / 1000}s; kernel restarted (state lost)`));
			}, timeoutMs);
			const onAbort = () => {
				this.pending.delete(id);
				clearTimeout(timer);
				this.dispose();
				reject(new Error("REPL execution aborted; kernel restarted (state lost)"));
			};
			signal?.addEventListener("abort", onAbort, { once: true });
			this.pending.set(id, {
				resolve: (r) => {
					clearTimeout(timer);
					signal?.removeEventListener("abort", onAbort);
					resolve(r);
				},
				reject: (e) => {
					clearTimeout(timer);
					signal?.removeEventListener("abort", onAbort);
					reject(e);
				},
			});
			this.proc.stdin?.write(`${JSON.stringify({ id, code })}\n`);
		});
	}

	dispose() {
		this.proc.kill("SIGKILL");
	}
}

function readNotes(dir: string): Array<{ name: string; content: string }> {
	if (!existsSync(dir)) return [];
	return readdirSync(dir)
		.filter((f) => f.endsWith(".md"))
		.sort()
		.map((f) => ({ name: f.slice(0, -3), content: readFileSync(join(dir, f), "utf8").trim() }));
}

function buildHarnessSection(cwd: string): string {
	const harnessDir = join(cwd, ".pi", "harness");
	const prompts = readNotes(join(harnessDir, "prompts"));
	const memories = readNotes(join(harnessDir, "memory"));
	if (prompts.length === 0 && memories.length === 0) return "";
	const parts = ["## Continual harness", "Self-managed notes (CRUD via the repl tool's `harness` object)."];
	if (prompts.length > 0) {
		parts.push("### Prompt notes");
		for (const p of prompts) parts.push(`#### ${p.name}\n${p.content}`);
	}
	if (memories.length > 0) {
		parts.push("### Memories");
		for (const m of memories) parts.push(`#### ${m.name}\n${m.content}`);
	}
	return parts.join("\n\n");
}

const REFINE_PROMPT = (trajectoryPath: string, focus: string) => `You are the /refine step of a self-improving agent harness. Read the trajectory of the agent's recent session at ${trajectoryPath} and inspect the current harness state under .pi/harness/ (memory/, prompts/, agents/) and .pi/skills/.

${focus ? `Focus on: ${focus}` : "Look for repeated failures worth remembering, or reusable tactics worth promoting."}

Apply the SMALLEST relevant change that improves future sessions: create or update ONE file — a memory note (.pi/harness/memory/<name>.md), a prompt note (.pi/harness/prompts/<name>.md), an agent spec (.pi/harness/agents/<name>.md), or a skill (.pi/skills/<name>/SKILL.md with YAML frontmatter containing name and description). Keep notes short and factual. Do not rewrite existing files wholesale; make targeted edits. Do not touch anything outside those directories.

Then append one JSON line to .pi/harness/refine-log.jsonl: {"ts": "<iso8601>", "trigger": "<what you observed>", "action": "<create|update|delete> <kind> <name>", "reason": "<why>"}

If the trajectory contains nothing worth persisting, change no files and append a log line with "action": "none". Finally, print a one-line summary.`;

export default function (pi: ExtensionAPI) {
	let kernel: Kernel | null = null;
	let refineProc: ChildProcess | null = null;
	const goal = { text: null as string | null, gate: null as string | null, turns: 0, maxTurns: 12 };

	pi.registerFlag("goal", { description: "Start with an autonomous goal (pi-continual)", type: "string" });

	pi.registerTool({
		name: "repl",
		label: "REPL",
		description: `Persistent Python REPL: state (variables, imports) survives across calls for the whole session. Output = captured stdout/stderr plus the repr of a trailing bare expression. Pre-loaded globals:
- rlm(task, name=None, agent=None, model=None) -> handle. Spawns a persistent sub-agent (a full pi session) and returns IMMEDIATELY. handle.wait(timeout=600) blocks and returns its final output; handle.running(); handle.result(). Calling rlm() again with the same name sends a follow-up turn into that same sub-agent session (it keeps its context). agent= names a spec created via harness ("agent" kind). Fan out parallel work by calling rlm() several times before waiting.
- subagents() -> dict of all sub-agent handles by name.
- harness: continual-harness CRUD. harness.create(kind, name, content, description="") / update / get / delete / list(kind=None). Kinds: "memory" and "prompt" (injected into your system prompt every turn), "agent" (sub-agent spec for rlm), "skill" (native pi skill; description required).
- history(n=None) -> last n entries of this session's JSONL (your own full history, including compacted-away context).`,
		promptSnippet:
			"Persistent Python REPL with rlm() sub-agents, continual-harness CRUD (memory/prompt/agent/skill), and history()",
		promptGuidelines: [
			"Use repl to fan out independent sub-tasks via rlm() in parallel instead of doing long serial work yourself.",
			"When you notice a repeated failure or a reusable tactic, persist it via repl: harness.create('memory'|'prompt'|'skill'|'agent', ...).",
		],
		parameters: Type.Object({
			code: Type.String({ description: "Python code to execute in the persistent kernel" }),
			timeoutSeconds: Type.Optional(Type.Number({ description: "Execution timeout (default 120). On timeout the kernel restarts and state is lost." })),
		}),
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			if (!kernel || !kernel.alive) {
				kernel = new Kernel(ctx.cwd, ctx.sessionManager.getSessionFile() ?? undefined);
			}
			const timeoutMs = Math.max(1, params.timeoutSeconds ?? 120) * 1000;
			const { ok, output } = await kernel.execute(params.code, timeoutMs, signal);
			if (!ok) throw new Error(truncate(output) || "REPL execution failed");
			return {
				content: [{ type: "text", text: truncate(output) || "(no output)" }],
				details: {},
			};
		},
	});

	pi.registerTool({
		name: "refine",
		label: "Refine",
		description:
			"Schedule background self-refinement: a separate agent reads this session's trajectory and applies the smallest useful edit to the harness state (.pi/harness, .pi/skills). Non-blocking. Call it when you notice a repeated failure or a reusable tactic.",
		parameters: Type.Object({
			focus: Type.Optional(Type.String({ description: "Optional specific observation to refine on" })),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const message = runRefine(params.focus ?? "", ctx);
			return { content: [{ type: "text", text: message }], details: {} };
		},
	});

	pi.registerTool({
		name: "goal_complete",
		label: "Goal complete",
		description:
			"Declare the current autonomous goal complete. If a gate command is configured, it runs first; a failing gate rejects completion and returns its output — fix the problems and try again.",
		parameters: Type.Object({
			summary: Type.String({ description: "What was accomplished and how it was verified" }),
		}),
		async execute(_toolCallId, params, signal) {
			if (!goal.text) {
				return { content: [{ type: "text", text: "No active goal." }], details: {} };
			}
			if (goal.gate) {
				const result = await pi.exec("bash", ["-lc", goal.gate], { signal, timeout: 300_000 });
				if (result.code !== 0) {
					throw new Error(
						`Gate command failed (exit ${result.code}): ${goal.gate}\n${truncate(`${result.stdout}\n${result.stderr}`, 8000)}`,
					);
				}
			}
			const finished = goal.text;
			goal.text = null;
			goal.turns = 0;
			return {
				content: [{ type: "text", text: `Goal complete${goal.gate ? " (gate passed)" : ""}: ${finished}\n${params.summary}` }],
				details: {},
				terminate: true,
			};
		},
	});

	function runRefine(focus: string, ctx: { cwd: string; sessionManager: { getEntries(): unknown[] }; ui: { notify(msg: string, level?: "info" | "warning" | "error"): void; setStatus(key: string, text: string | undefined): void } }): string {
		if (refineProc) return "A refinement is already running; wait for it to finish.";
		const harnessDir = join(ctx.cwd, ".pi", "harness");
		mkdirSync(join(harnessDir, "refine"), { recursive: true });

		const lines: string[] = [];
		const entries = ctx.sessionManager.getEntries().slice(-120) as Array<{ type: string; message?: { role: string; content: unknown } }>;
		for (const entry of entries) {
			if (entry.type !== "message" || !entry.message) continue;
			const { role, content } = entry.message;
			const text =
				typeof content === "string"
					? content
					: (content as Array<{ type: string; text?: string; name?: string; arguments?: unknown }>)
							.map((b) => b.text ?? (b.name ? `[tool call: ${b.name} ${JSON.stringify(b.arguments ?? {})}]` : ""))
							.filter(Boolean)
							.join("\n");
			if (text) lines.push(`--- ${role} ---\n${truncate(text, 3000)}`);
		}
		if (lines.length === 0) return "Nothing in the trajectory yet; nothing to refine.";

		const trajectoryPath = join(harnessDir, "refine", "trajectory.txt");
		writeFileSync(trajectoryPath, lines.join("\n\n"), "utf8");

		const logPath = join(harnessDir, "refine", "last-run.txt");
		refineProc = spawn("pi", ["-p", "--no-extensions", "-nc", REFINE_PROMPT(trajectoryPath, focus)], {
			cwd: ctx.cwd,
			stdio: ["ignore", "pipe", "pipe"],
		});
		const chunks: Buffer[] = [];
		refineProc.stdout?.on("data", (c: Buffer) => chunks.push(c));
		refineProc.stderr?.on("data", (c: Buffer) => chunks.push(c));
		ctx.ui.setStatus("continual", "refine: running…");
		refineProc.on("exit", (code) => {
			refineProc = null;
			const output = Buffer.concat(chunks).toString("utf8").trim();
			writeFileSync(logPath, output, "utf8");
			ctx.ui.setStatus("continual", undefined);
			const summary = output.split("\n").filter(Boolean).pop() ?? "(no output)";
			ctx.ui.notify(code === 0 ? `refine: ${truncate(summary, 200)}` : `refine failed (exit ${code}), see ${logPath}`, code === 0 ? "info" : "error");
		});
		return `Refinement scheduled in the background${focus ? ` (focus: ${focus})` : ""}. Results land in .pi/harness and .pi/harness/refine-log.jsonl.`;
	}

	pi.on("before_agent_start", async (event, ctx) => {
		const section = buildHarnessSection(ctx.cwd);
		let systemPrompt = event.systemPrompt;
		if (section) systemPrompt += `\n\n${section}`;
		if (goal.text) {
			systemPrompt += `\n\n## Autonomous goal\nPersistent objective: ${goal.text}\nKeep working toward it each turn. When it is genuinely done and verified, call the goal_complete tool.${goal.gate ? ` Completion is gated by: ${goal.gate}` : ""}`;
		}
		return section || goal.text ? { systemPrompt } : undefined;
	});

	pi.on("agent_settled", async (_event, ctx) => {
		if (!goal.text) return;
		if (goal.turns >= goal.maxTurns) {
			const abandoned = goal.text;
			goal.text = null;
			goal.turns = 0;
			ctx.ui.notify(`goal: turn limit (${goal.maxTurns}) reached, stopping: ${abandoned}`, "warning");
			return;
		}
		goal.turns++;
		pi.sendUserMessage(
			`[pi-continual autonomous mode, turn ${goal.turns}/${goal.maxTurns}] The goal is not yet complete: ${goal.text}\nContinue working toward it. If it is done and verified, call goal_complete.`,
		);
	});

	pi.on("session_start", async (_event, ctx) => {
		const flagGoal = pi.getFlag("goal");
		if (typeof flagGoal === "string" && flagGoal.trim()) {
			goal.text = flagGoal.trim();
			goal.turns = 0;
			// In -p mode the CLI prompt starts the turn; interactively nothing would.
			if (ctx.hasUI && ctx.isIdle()) {
				pi.sendUserMessage(`Work toward this goal: ${goal.text}\nWhen it is done and verified, call goal_complete.`);
			}
		}
	});

	pi.on("session_shutdown", async () => {
		kernel?.dispose();
		kernel = null;
	});

	pi.registerCommand("goal", {
		description: "Autonomous goal: /goal <text> to set, /goal to show, /goal off to clear",
		handler: async (args, ctx) => {
			const text = (args ?? "").trim();
			if (text === "off") {
				goal.text = null;
				goal.turns = 0;
				ctx.ui.notify("goal cleared", "info");
			} else if (text) {
				goal.text = text;
				goal.turns = 0;
				ctx.ui.notify(`goal set (max ${goal.maxTurns} turns): ${text}`, "info");
				pi.sendUserMessage(`Work toward this goal: ${text}\nWhen it is done and verified, call goal_complete.`);
			} else {
				ctx.ui.notify(goal.text ? `goal (turn ${goal.turns}/${goal.maxTurns}): ${goal.text}` : "no active goal", "info");
			}
		},
	});

	pi.registerCommand("gate", {
		description: "Goal gate command: /gate <cmd> to set, /gate to show, /gate off to clear",
		handler: async (args, ctx) => {
			const text = (args ?? "").trim();
			if (text === "off") {
				goal.gate = null;
				ctx.ui.notify("gate cleared", "info");
			} else if (text) {
				goal.gate = text;
				ctx.ui.notify(`gate set: ${text}`, "info");
			} else {
				ctx.ui.notify(goal.gate ? `gate: ${goal.gate}` : "no gate configured", "info");
			}
		},
	});

	pi.registerCommand("refine", {
		description: "Background self-refinement of the harness state: /refine [focus]",
		handler: async (args, ctx) => {
			ctx.ui.notify(runRefine((args ?? "").trim(), ctx), "info");
		},
	});
}