/**
 * Git Worktree Extension
 *
 * Spin up an AI agent in a git worktree to implement tasks in the background
 * while you continue working in the main branch.
 *
 * Commands:
 *   /worktree <branch> <task>  — create worktree + start agent
 *   /worktrees                 — list all worktree agents and their status
 *   /worktree-log <branch>     — show recent log output (and log file path)
 *   /worktree-done <branch>    — remove a finished worktree
 *
 * The worktree is placed as a sibling of your repo:
 *   <repoRoot>/../<repoName>-<branch>
 *
 * Agent output is streamed to:
 *   /tmp/pi-worktrees/<branch>.log
 *
 * Follow live:  !tail -f /tmp/pi-worktrees/<branch>.log
 *
 * DESIGN NOTES
 * ─────────────
 * Context:   The full conversation history is serialized and prepended to the
 *            task prompt so the agent sees the plan (or whatever was discussed).
 *
 * Visibility: Spawns pi in --mode json so we can parse events. The widget above
 *             the editor shows ⏳/✓/✗ per agent, elapsed time, and the last few
 *             tool calls / text snippets the agent produced.
 *
 * Questions:  The agent runs in -p (batch) mode and cannot pause to ask you
 *             anything. Give it enough context upfront (the conversation carry-
 *             over handles this). If the agent gets stuck it will make its best
 *             guess or fail — check the log and use /worktree-done to clean up.
 */

import { execSync } from "node:child_process";
import { spawn } from "node:child_process";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
	SessionEntry,
} from "@mariozechner/pi-coding-agent";
import { convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";

// ─── Types ────────────────────────────────────────────────────────────────────

interface WorktreeAgent {
	branch: string;
	worktreePath: string;
	task: string;
	logFile: string;
	startTime: number;
	status: "running" | "done" | "error";
	exitCode?: number;
	/** Last few activity lines parsed from the JSON event stream */
	recentActivity: string[];
	/** Accumulated final-text output from all assistant turns */
	finalOutput: string;
}

type Ctx = ExtensionContext | ExtensionCommandContext;

// ─── Helpers ──────────────────────────────────────────────────────────────────

function shortenPath(p: string): string {
	const home = os.homedir();
	return p.startsWith(home) ? `~${p.slice(home.length)}` : p;
}

function formatToolCall(toolName: string, args: Record<string, unknown>): string {
	switch (toolName) {
		case "bash": {
			const cmd = ((args.command as string) || "").split("\n")[0].slice(0, 60);
			return `$ ${cmd}`;
		}
		case "read":
			return `read ${shortenPath((args.file_path ?? args.path ?? "?") as string)}`;
		case "write":
			return `write ${shortenPath((args.file_path ?? args.path ?? "?") as string)}`;
		case "edit":
			return `edit ${shortenPath((args.file_path ?? args.path ?? "?") as string)}`;
		case "grep":
			return `grep /${args.pattern}/ in ${shortenPath((args.path ?? ".") as string)}`;
		case "find":
			return `find ${args.pattern} in ${shortenPath((args.path ?? ".") as string)}`;
		default:
			return `${toolName}(${JSON.stringify(args).slice(0, 40)})`;
	}
}

// ─── Extension ────────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
	const agents = new Map<string, WorktreeAgent>();

	// ── Git helpers ────────────────────────────────────────────────────────────

	function getRepoRoot(cwd: string): string | null {
		try {
			return execSync(`git -C "${cwd}" rev-parse --show-toplevel`, { encoding: "utf-8" }).trim();
		} catch {
			return null;
		}
	}

	function branchExists(repoRoot: string, branch: string): boolean {
		try {
			execSync(`git -C "${repoRoot}" rev-parse --verify "${branch}"`, { stdio: "ignore" });
			return true;
		} catch {
			return false;
		}
	}

	// ── Widget ─────────────────────────────────────────────────────────────────

	function refreshWidget(ctx: Ctx) {
		const lines: string[] = [];
		for (const wt of agents.values()) {
			const icon = wt.status === "running" ? "⏳" : wt.status === "done" ? "✓" : "✗";
			const elapsed = Date.now() - wt.startTime;
			const timeStr = elapsed < 60_000 ? `${Math.round(elapsed / 1000)}s` : `${Math.round(elapsed / 60_000)}m`;
			lines.push(`${icon} ${wt.branch} [${timeStr}]`);
			for (const line of wt.recentActivity.slice(-3)) {
				lines.push(`    ${line}`);
			}
		}
		ctx.ui.setWidget("worktrees", lines);
	}

	pi.on("turn_start", async (_e, ctx) => refreshWidget(ctx));
	pi.on("agent_end", async (_e, ctx) => refreshWidget(ctx));

	// ── /worktree <branch> <task> ──────────────────────────────────────────────

	pi.registerCommand("worktree", {
		description: "Create a git worktree and run an agent in it: /worktree <branch> <task>",
		handler: async (args, ctx) => {
			const trimmed = args.trim();
			const spaceIdx = trimmed.indexOf(" ");
			if (spaceIdx === -1) {
				ctx.ui.notify("Usage: /worktree <branch-name> <task description>", "error");
				return;
			}

			const branch = trimmed.slice(0, spaceIdx);
			const task = trimmed.slice(spaceIdx + 1).trim();
			if (!task) {
				ctx.ui.notify("Please provide a task description after the branch name", "error");
				return;
			}

			if (agents.get(branch)?.status === "running") {
				ctx.ui.notify(`An agent for branch '${branch}' is already running`, "error");
				return;
			}

			const repoRoot = getRepoRoot(ctx.cwd);
			if (!repoRoot) {
				ctx.ui.notify("Not inside a git repository", "error");
				return;
			}

			// ── Serialize current conversation as context ──────────────────────
			// The agent needs to see whatever was discussed (e.g. the plan) so it
			// can act on "implement the above plan" style instructions.
			const branch_ = ctx.sessionManager.getBranch();
			const sessionMessages = branch_
				.filter((e): e is SessionEntry & { type: "message" } => e.type === "message")
				.map((e) => e.message);

			let conversationContext = "";
			if (sessionMessages.length > 0) {
				try {
					const llmMessages = convertToLlm(sessionMessages);
					conversationContext = serializeConversation(llmMessages);
				} catch {
					// If serialization fails, proceed without context
				}
			}

			// ── Create worktree ────────────────────────────────────────────────
			const repoName = path.basename(repoRoot);
			const safeBranch = branch.replace(/[^\w.-]/g, "_");
			const worktreePath = path.resolve(repoRoot, "..", `${repoName}-${safeBranch}`);

			try {
				if (branchExists(repoRoot, branch)) {
					execSync(`git -C "${repoRoot}" worktree add "${worktreePath}" "${branch}"`, { stdio: "pipe" });
				} else {
					execSync(`git -C "${repoRoot}" worktree add "${worktreePath}" -b "${branch}"`, { stdio: "pipe" });
				}
			} catch (e: any) {
				if (!fs.existsSync(path.join(worktreePath, ".git"))) {
					const msg = (e as any).stderr?.toString().trim() || (e as Error).message;
					ctx.ui.notify(`Failed to create worktree: ${msg}`, "error");
					return;
				}
				// Already registered — continue
			}

			// ── Log + temp files ───────────────────────────────────────────────
			const logDir = path.join(os.tmpdir(), "pi-worktrees");
			fs.mkdirSync(logDir, { recursive: true });
			const logFile = path.join(logDir, `${safeBranch}.log`);

			// Build the full prompt the agent will receive.
			// If there is prior conversation, prepend it so the agent has all
			// the context it needs (e.g. the plan written in the last message).
			const fullPrompt =
				conversationContext.trim()
					? [
							"The following is the conversation that led to this task. Use it as context.",
							"",
							"<conversation>",
							conversationContext.trim(),
							"</conversation>",
							"",
							"Your task:",
							task,
						].join("\n")
					: task;

			// Temp file for the prompt (avoids shell-quoting issues with long text)
			const promptFile = path.join(logDir, `${safeBranch}-prompt.md`);
			fs.writeFileSync(promptFile, fullPrompt);

			// System-prompt addendum: ground the agent in its worktree context
			const systemFile = path.join(logDir, `${safeBranch}-system.md`);
			fs.writeFileSync(
				systemFile,
				[
					`You are an AI coding agent working in a git worktree for branch '${branch}'.`,
					`Working directory: ${worktreePath}`,
					`Main repository:   ${repoRoot}`,
					``,
					`Implement the requested changes fully and correctly.`,
					`When you are finished, stage all changed files and create a git commit with a descriptive message.`,
					`You cannot ask the user questions — work autonomously with the context provided.`,
				].join("\n"),
			);

			// ── Register + spawn ───────────────────────────────────────────────
			const agent: WorktreeAgent = {
				branch,
				worktreePath,
				task,
				logFile,
				startTime: Date.now(),
				status: "running",
				recentActivity: [],
				finalOutput: "",
			};
			agents.set(branch, agent);

			const logStream = fs.createWriteStream(logFile);
			logStream.write(
				[
					`=== pi-worktree agent started ===`,
					`Branch:   ${branch}`,
					`Path:     ${worktreePath}`,
					`Task:     ${task}`,
					`Started:  ${new Date().toISOString()}`,
					`Context:  ${sessionMessages.length} messages from current session`,
					`Log:      tail -f ${logFile}`,
					``,
					``,
				].join("\n"),
			);

			// Use --mode json so we can parse events and populate the widget.
			// The prompt is passed via @file reference to avoid shell-length limits.
			const proc = spawn(
				"pi",
				[
					"--mode", "json",
					"-p",
					"--no-session",
					"--append-system-prompt", systemFile,
					`@${promptFile}`,
				],
				{ cwd: worktreePath, stdio: ["ignore", "pipe", "pipe"] },
			);

			// Parse JSON event stream for widget updates; also mirror to log file
			let buffer = "";
			proc.stdout?.on("data", (chunk: Buffer) => {
				const text = chunk.toString();
				logStream.write(text);
				buffer += text;
				const lines = buffer.split("\n");
				buffer = lines.pop() ?? "";
				for (const line of lines) {
					if (!line.trim()) continue;
					try {
						const event = JSON.parse(line) as Record<string, unknown>;
						handleJsonEvent(agent, event);
					} catch {
						// Non-JSON line — ignore for widget purposes
					}
				}
			});

			proc.stderr?.on("data", (d: Buffer) => {
				logStream.write(`[stderr] ${d}`);
			});

			proc.on("close", (code) => {
				const wt = agents.get(branch);
				if (wt) {
					wt.status = code === 0 ? "done" : "error";
					wt.exitCode = code ?? undefined;
					wt.recentActivity.push(code === 0 ? "✓ done" : `✗ exited ${code}`);
				}
				logStream.write(`\n=== Agent exited with code ${code} — ${new Date().toISOString()} ===\n`);
				logStream.end();
				cleanup();

				// ── Hand results back to the main session ──────────────────────
				// Inject a message into the main conversation so you (and the main
				// agent) can see what the worktree agent did and discuss it.
				const elapsed = wt ? Math.round((Date.now() - wt.startTime) / 1000) : 0;
				const summary = wt?.finalOutput.trim() || "(no text output)";
				const statusLine = code === 0
					? `✓ Worktree agent for **${branch}** finished in ${elapsed}s`
					: `✗ Worktree agent for **${branch}** failed (exit ${code}) after ${elapsed}s`;

				pi.sendMessage(
					{
						customType: "worktree-result",
						content: [
							statusLine,
							`Worktree: \`${worktreePath}\``,
							`Task: ${task}`,
							``,
							`**Agent output:**`,
							summary,
						].join("\n"),
						display: true,
					},
					{
						// Queue for the next time the user sends a message.
						// Change to "steer" + triggerTurn:true if you want an
						// immediate LLM response the moment the agent finishes.
						deliverAs: "nextTurn",
					},
				);
			});

			proc.on("error", (err) => {
				const wt = agents.get(branch);
				if (wt) {
					wt.status = "error";
					wt.recentActivity.push(`spawn error: ${err.message}`);
				}
				logStream.write(`\n[spawn error] ${err.message}\n`);
				logStream.end();
				cleanup();
			});

			function cleanup() {
				try { fs.unlinkSync(systemFile); } catch {}
				try { fs.unlinkSync(promptFile); } catch {}
			}

			ctx.ui.notify(`🌿 Worktree agent started for '${branch}'  —  tail -f ${logFile}`, "info");
			refreshWidget(ctx);
		},
	});

	// ── /worktrees ─────────────────────────────────────────────────────────────

	pi.registerCommand("worktrees", {
		description: "List all worktree agents and their status",
		handler: async (_args, ctx) => {
			if (agents.size === 0) {
				ctx.ui.notify("No worktree agents active this session", "info");
				return;
			}
			for (const wt of agents.values()) {
				const icon = wt.status === "running" ? "⏳" : wt.status === "done" ? "✓" : "✗";
				const elapsed = Date.now() - wt.startTime;
				const timeStr = elapsed < 60_000 ? `${Math.round(elapsed / 1000)}s` : `${Math.round(elapsed / 60_000)}m`;
				const last = wt.recentActivity.at(-1) ?? "(no activity yet)";
				ctx.ui.notify(`${icon} ${wt.branch} [${timeStr}]\n  task: ${wt.task}\n  last: ${last}`, "info");
			}
		},
	});

	// ── /worktree-log <branch> ─────────────────────────────────────────────────

	pi.registerCommand("worktree-log", {
		description: "Show recent log for a worktree agent: /worktree-log <branch>",
		handler: async (args, ctx) => {
			const branch = args.trim();
			const wt = agents.get(branch);
			if (!wt) {
				ctx.ui.notify(
					`No worktree agent for '${branch}'\nKnown: ${[...agents.keys()].join(", ") || "none"}`,
					"error",
				);
				return;
			}

			ctx.ui.notify(`Log: ${wt.logFile}\nLive follow:  !tail -f ${wt.logFile}`, "info");

			try {
				const raw = fs.readFileSync(wt.logFile, "utf-8");
				// Log is JSON-mode output — skip JSON lines, show text content only
				const readable = raw
					.split("\n")
					.filter((l) => {
						if (!l.trim()) return false;
						try { JSON.parse(l); return false; } catch { return true; }
					})
					.slice(-30)
					.join("\n");
				if (readable) ctx.ui.notify(`[${branch} — last readable lines]\n\n${readable}`, "info");
			} catch {
				// Log not readable yet
			}
		},
	});

	// ── /worktree-done <branch> ────────────────────────────────────────────────

	pi.registerCommand("worktree-done", {
		description: "Remove a finished worktree: /worktree-done <branch>",
		handler: async (args, ctx) => {
			const branch = args.trim();
			const wt = agents.get(branch);
			if (!wt) {
				ctx.ui.notify(`No worktree agent for '${branch}'`, "error");
				return;
			}

			if (wt.status === "running") {
				const ok = await ctx.ui.confirm(
					"Agent still running",
					`Force-remove the worktree for '${branch}' even though the agent has not finished?`,
				);
				if (!ok) return;
			}

			const repoRoot = getRepoRoot(ctx.cwd);
			if (repoRoot) {
				try {
					execSync(`git -C "${repoRoot}" worktree remove "${wt.worktreePath}" --force`, { stdio: "pipe" });
				} catch (e: any) {
					const msg = (e as any).stderr?.toString().trim() || (e as Error).message;
					ctx.ui.notify(`Warning: could not remove worktree directory: ${msg}`, "warning");
				}
			}

			agents.delete(branch);
			refreshWidget(ctx);
			ctx.ui.notify(`Removed worktree '${branch}'`, "info");
		},
	});

	// ── JSON event parser ──────────────────────────────────────────────────────
	// Parses --mode json events to populate recentActivity for the widget.

	function handleJsonEvent(agent: WorktreeAgent, event: Record<string, unknown>) {
		// Tool starting — show what it's about to do
		if (event.type === "tool_execution_start") {
			const name = event.toolName as string;
			const rawArgs = event.args as Record<string, unknown> | undefined;
			if (name && rawArgs) {
				pushActivity(agent, `→ ${formatToolCall(name, rawArgs)}`);
			}
			return;
		}

		// Assistant text — grab a brief preview for the widget and accumulate
		// the full text so we can hand it back to the main session on exit.
		if (event.type === "message_end") {
			const msg = event.message as { role?: string; content?: unknown[] } | undefined;
			if (msg?.role === "assistant" && Array.isArray(msg.content)) {
				for (const part of msg.content) {
					const p = part as { type?: string; text?: string };
					if (p.type === "text" && p.text) {
						// Widget: short preview
						const preview = p.text.replace(/\s+/g, " ").slice(0, 70);
						pushActivity(agent, `💬 ${preview}${p.text.length > 70 ? "…" : ""}`);
						// Handback: accumulate full text across all turns
						agent.finalOutput += (agent.finalOutput ? "\n\n" : "") + p.text;
						break;
					}
				}
			}
		}
	}

	function pushActivity(agent: WorktreeAgent, line: string) {
		agent.recentActivity.push(line);
		// Keep a rolling window so memory doesn't grow unboundedly
		if (agent.recentActivity.length > 50) {
			agent.recentActivity.splice(0, agent.recentActivity.length - 50);
		}
	}
}