/**
 * Pi Loop Extension
 *
 * A reusable Pi extension that adds a supervised loop harness around the
 * normal Pi coding agent. The extension owns workflow state, node order,
 * tool scope, steering, validation, approval gates, and completion.
 *
 * This extension ships with NO predefined loops so you can build your own
 * workflows on top of the generic runtime.
 *
 * See src/loops/ to add your own loop definitions.
 * See skills/pi-loops/SKILL.md for the agent-facing guide.
 */

import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import type { LoopDefinition, LoopReport, RunState } from "./types.ts";
import { registerDefaultCodeHandlers } from "./nodes.ts";
import { registerLoopPackages } from "./loops/index.ts";
import { LoopRuntime } from "./runtime.ts";
import { isRunActive, restoreRunState } from "./state.ts";
import { registerLoopReportTool } from "./tools.ts";
import { renderLoopWidget } from "./ui.ts";

export default function piLoopExtension(pi: ExtensionAPI): void {
  // ---------------------------------------------------------------------------
  // State
  // ---------------------------------------------------------------------------

  /** Registered loop definitions (id -> definition) */
  const loopDefinitions = new Map<string, LoopDefinition>();

  /** Currently active loop runtimes (runId -> runtime) */
  const activeRuntimes = new Map<string, LoopRuntime>();

  /** The most recent/active run ID */
  let activeRunId: string | null = null;
  const completionNotified = new Set<string>();

  /** Currently active runtime for the current session */
  let currentRuntime: LoopRuntime | null = null;

  // ---------------------------------------------------------------------------
  // Initialization: Register loop definitions
  // ---------------------------------------------------------------------------

  function registerLoopDefinition(def: LoopDefinition): void {
    loopDefinitions.set(def.id, def);
  }

  registerDefaultCodeHandlers();
  registerLoopPackages(registerLoopDefinition);

  // ---------------------------------------------------------------------------
  // Loop Report Tool
  // ---------------------------------------------------------------------------

  registerLoopReportTool(pi, (report: LoopReport) => {
    if (currentRuntime) {
      currentRuntime.handleReport(report);
    }
  });

  // ---------------------------------------------------------------------------
  // Tool Call Observer + Policy Enforcement
  // ---------------------------------------------------------------------------

  pi.on("tool_call", (event, _ctx) => {
    if (!currentRuntime) return;
    const policyBlock = currentRuntime.enforceToolPolicy(event.toolName, event.input);
    if (policyBlock) return policyBlock;
    currentRuntime.recordToolCallEvidence(event.toolName, event.input);
  });

  pi.on("tool_result", (event, _ctx) => {
    if (!currentRuntime) return;
    if (event.isError) {
      currentRuntime.recordToolCallEvidence(event.toolName, event.input, true);
    }
  });

  // ---------------------------------------------------------------------------
  // Context cleaning — preserves current LLM node history while discarding
  // stale messages from previous nodes in the same run.
  // ---------------------------------------------------------------------------

  pi.on("context", async (event) => {
    if (!currentRuntime?.shouldCleanContextForActiveNode()) return;

    const markerIndex = findLatestLoopNodeContextMarker(event.messages as unknown as Array<Record<string, unknown>>);
    if (markerIndex < 0) return;

    return { messages: event.messages.slice(markerIndex) };
  });

  // ---------------------------------------------------------------------------
  // Session Management
  // ---------------------------------------------------------------------------

  pi.on("session_start", async (_event, ctx) => {
    const restoredRun = restoreRunState(ctx);
    if (restoredRun && isRunActive(restoredRun)) {
      const def = loopDefinitions.get(restoredRun.loopId);
      if (def) {
        const runtime = new LoopRuntime({ pi, ctx }, def, restoredRun);
        activeRuntimes.set(restoredRun.id, runtime);
        activeRunId = restoredRun.id;
        currentRuntime = runtime;

        ctx.ui.notify(`Restored active loop: ${restoredRun.loopId} (${restoredRun.status})`, "info");
        ctx.ui.setStatus("pi-loop", `Loop: ${restoredRun.loopId} | Status: ${restoredRun.status}`);
      }
    }
  });

  pi.on("session_shutdown", async () => {
    currentRuntime = null;
    activeRunId = null;
  });

  // ---------------------------------------------------------------------------
  // Commands
  // ---------------------------------------------------------------------------

  pi.registerCommand("loop", {
    description: "Manage loop runs. Usage: /loop start <def> [input...] | /loop status | /loop list | /loop cancel | /loop pause | /loop resume",
    handler: async (args: string, ctx: ExtensionCommandContext) => {
      const parts = args.trim().split(/\s+/);
      const subcommand = parts[0]?.toLowerCase();

      if (!subcommand) {
        ctx.ui.notify(
          "Loop commands:\n  /loop start <definition> [input...] - Start a new loop\n  /loop status - Show current loop status\n  /loop list - List available loop definitions\n  /loop cancel - Cancel current loop\n  /loop pause - Pause current loop\n  /loop resume - Resume paused loop",
          "info",
        );
        return;
      }

      switch (subcommand) {
        case "start":
          await handleStart(parts.slice(1), ctx);
          break;
        case "status":
          handleStatus(ctx);
          break;
        case "list":
        case "definitions":
          handleListDefinitions(ctx);
          break;
        case "cancel":
          handleCancel(ctx);
          break;
        case "pause":
          handlePause(ctx);
          break;
        case "resume":
          handleResume(ctx);
          break;
        default:
          ctx.ui.notify(`Unknown loop subcommand: ${subcommand}. Try: start, status, list, cancel, pause, resume`, "error");
      }
    },
  });

  // ---------------------------------------------------------------------------
  // Command Handlers
  // ---------------------------------------------------------------------------

  async function handleStart(args: string[], ctx: ExtensionCommandContext): Promise<void> {
    if (args.length === 0) {
      ctx.ui.notify(
        "Available definitions:\n" +
        Array.from(loopDefinitions.keys()).map((id) => `  ${id}: ${loopDefinitions.get(id)?.label ?? id}`).join("\n"),
        "info",
      );
      return;
    }

    const defName = args[0];
    const def = loopDefinitions.get(defName);

    if (!def) {
      ctx.ui.notify(`Loop definition "${defName}" not found. Use /loop list to see available definitions.`, "error");
      return;
    }

    const input: Record<string, string> = {};
    const positionalFields = def.trigger.inputShape?.map((field) => field.name) ?? [];
    const inputArgs = args.slice(1);
    let positionalIndex = 0;
    for (let i = 0; i < inputArgs.length; i++) {
      const kv = inputArgs[i].split("=");
      if (kv.length === 2 && kv[0]) {
        input[kv[0]] = kv.slice(1).join("=");
      } else {
        const field = positionalFields[positionalIndex] ?? `arg${i}`;
        input[field] = inputArgs[i];
        input[`arg${i}`] = inputArgs[i];
        positionalIndex++;
      }
    }

    ctx.ui.notify(`Starting loop: ${def.label}`, "info");

    const runtime = new LoopRuntime({ pi, ctx }, def, undefined, input);

    runtime.onStateChangeCallback((run) => {
      const statusLine = `Loop: ${run.loopId}#${run.id.slice(0, 8)} | Node: ${getCurrentNodeLabel(run)} | Attempt: ${run.currentAttemptNumber} | Status: ${run.status}`;
      ctx.ui.setStatus("pi-loop", statusLine);

      const lines = renderLoopWidget(run);
      ctx.ui.setWidget("pi-loop", lines);

      if (run.status === "completed") {
        if (!completionNotified.has(run.id)) {
          completionNotified.add(run.id);
          const summaryLines = buildCompletionSummary(run);
          ctx.ui.notify(`🎉 Loop completed: ${run.loopId}\n${summaryLines.join("\n")}`, "info");
          ctx.ui.setWidget("pi-loop", [...lines, "", ...summaryLines]);
        }
      } else if (run.status === "failed") {
        ctx.ui.notify(`❌ Loop failed: ${run.loopId}`, "error");
      }
    });

    activeRuntimes.set(runtime.getRun().id, runtime);
    activeRunId = runtime.getRun().id;
    currentRuntime = runtime;

    const initialChoice = await ctx.ui.select("Run mode:", [
      "Run now",
      "Run in background",
      "Step through nodes manually",
    ]);

    if (initialChoice === "Run now") {
      await runtime.start();
    } else if (initialChoice === "Step through nodes manually") {
      ctx.ui.notify("Step mode: use /loop status and /loop resume to continue", "info");
    } else {
      ctx.ui.notify("Loop started in background", "info");
      runtime.start().catch((err) => {
        ctx.ui.notify(`Loop error: ${err instanceof Error ? err.message : String(err)}`, "error");
      });
    }
  }

  function handleStatus(ctx: ExtensionCommandContext): void {
    if (!currentRuntime) {
      ctx.ui.notify("No active loop. Start one with /loop start <definition>", "info");
      return;
    }

    const run = currentRuntime.getRun();
    const lines = [
      `Loop: ${run.loopId} (${run.id.slice(0, 8)})`,
      `Status: ${run.status}`,
      `Current node: ${getCurrentNodeLabel(run)} (index ${run.currentNodeIndex})`,
      `Attempt: ${run.currentAttemptNumber}`,
      `Nodes: ${run.nodes.filter((n) => n.status === "completed").length}/${run.nodes.length} completed`,
      `Turns: ${run.totalTurns}`,
      `Tool calls: ${run.totalToolCalls}`,
      `Errors: ${run.evidence.errorCount}`,
      `Duration: ${formatDuration(Date.now() - run.startedAt)}`,
      ``,
      `Nodes:`,
      ...run.nodes.map((n, i) => {
        const current = i === run.currentNodeIndex ? " ←" : "";
        const summary = n.result?.summary ? ` — ${n.result.summary}` : n.gateResult?.feedback ? ` — ${n.gateResult.feedback}` : "";
        return `  [${n.status}] ${n.nodeDef.label}${current}${summary}`;
      }),
    ];

    ctx.ui.notify(lines.join("\n"), "info");
  }

  function handleListDefinitions(ctx: ExtensionCommandContext): void {
    const lines = Array.from(loopDefinitions.values()).map(
      (def) => `  ${def.id} (v${def.version}): ${def.label} - ${def.description}`,
    );
    ctx.ui.notify(`Available loop definitions:\n${lines.join("\n")}`, "info");
  }

  function handleCancel(ctx: ExtensionCommandContext): void {
    if (!currentRuntime) {
      ctx.ui.notify("No active loop to cancel.", "info");
      return;
    }
    currentRuntime.cancel();
    ctx.ui.notify("Loop cancelled.", "info");
    ctx.ui.setStatus("pi-loop", undefined);
    ctx.ui.setWidget("pi-loop", undefined);
    currentRuntime = null;
  }

  function handlePause(ctx: ExtensionCommandContext): void {
    if (!currentRuntime) {
      ctx.ui.notify("No active loop to pause.", "info");
      return;
    }
    currentRuntime.pause();
    ctx.ui.notify("Loop paused. Resume with /loop resume", "info");
  }

  function handleResume(ctx: ExtensionCommandContext): void {
    if (!currentRuntime) {
      const restoredRun = restoreRunState(ctx);
      if (restoredRun && restoredRun.status === "paused") {
        const def = loopDefinitions.get(restoredRun.loopId);
        if (def) {
          const runtime = new LoopRuntime({ pi, ctx }, def, restoredRun);
          activeRuntimes.set(restoredRun.id, runtime);
          activeRunId = restoredRun.id;
          currentRuntime = runtime;
          runtime.resume();
          ctx.ui.notify("Resumed loop from saved state.", "info");
          return;
        }
      }
      ctx.ui.notify("No active loop to resume.", "info");
      return;
    }
    currentRuntime.resume();
    ctx.ui.notify("Loop resumed.", "info");
  }

  // ---------------------------------------------------------------------------
  // Helpers
  // ---------------------------------------------------------------------------

  function getCurrentNodeLabel(run: RunState): string {
    const node = run.nodes[run.currentNodeIndex];
    return node ? `${node.nodeDef.label} (${node.status})` : "none";
  }

  function formatDuration(ms: number): string {
    const seconds = Math.floor(ms / 1000);
    const minutes = Math.floor(seconds / 60);
    const hours = Math.floor(minutes / 60);
    if (hours > 0) return `${hours}h ${minutes % 60}m`;
    if (minutes > 0) return `${minutes}m ${seconds % 60}s`;
    return `${seconds}s`;
  }

  function buildCompletionSummary(run: RunState): string[] {
    const metrics = readRunMetrics(run);
    const manifest = (run.artifacts.manifest as any) ?? {};
    const counts = manifest.counts ?? {};
    const tokens = metrics?.totals?.allTokens ?? metrics?.totals?.childCallTokens;
    const lines: string[] = [];
    if (metrics?.durationMs) lines.push(`Duration: ${formatDuration(Number(metrics.durationMs))}`);
    if (counts.facts != null) {
      lines.push(`Facts: ${counts.facts} | Req: ${counts.requirements ?? 0} | Info: ${counts.info ?? 0} | L3: ${counts.l3 ?? 0}`);
    }
    if (tokens) {
      lines.push(`Tokens: ↑${formatNumber(tokens.input)} ↓${formatNumber(tokens.output)} total ${formatNumber(tokens.totalTokens)}`);
      if (Number(tokens.cacheRead ?? 0) || Number(tokens.cacheWrite ?? 0)) {
        lines.push(`Cache: R${formatNumber(tokens.cacheRead)} W${formatNumber(tokens.cacheWrite)}`);
      }
      lines.push(`Cost reported: $${Number(tokens.cost ?? 0).toFixed(6)}`);
    }
    const stepLines = buildStepDurationLines(metrics);
    if (stepLines.length) lines.push("", "Main step durations:", ...stepLines);
    const out = manifest.outputs?.metrics ?? (run.artifacts.outputFolder ? `${run.artifacts.outputFolder}/run-metrics.json` : undefined);
    if (out) lines.push(`Metrics: ${out}`);
    return lines.length ? lines : [`Run: ${run.id}`];
  }

  function readRunMetrics(run: RunState): any | undefined {
    const runDir = typeof run.artifacts.runDir === "string" ? run.artifacts.runDir : undefined;
    if (!runDir) return undefined;
    const metricsPath = join(runDir, "run-metrics.json");
    if (!existsSync(metricsPath)) return undefined;
    try {
      return JSON.parse(readFileSync(metricsPath, "utf8"));
    } catch {
      return undefined;
    }
  }

  function buildStepDurationLines(metrics: any | undefined): string[] {
    const steps = Array.isArray(metrics?.steps) ? metrics.steps : [];
    return steps
      .filter((step: any) => step?.status !== "pending")
      .map((step: any) => `${String(step.id).padEnd(34)} ${formatDurationPrecise(Number(step.durationMs ?? 0))}`);
  }

  function formatDurationPrecise(ms: number): string {
    const seconds = ms / 1000;
    if (seconds < 60) return `${seconds.toFixed(1)}s`;
    const minutes = Math.floor(seconds / 60);
    return `${minutes}m ${(seconds - minutes * 60).toFixed(1)}s`;
  }

  function formatNumber(value: unknown): string {
    return Number(value ?? 0).toLocaleString();
  }
}

function findLatestLoopNodeContextMarker(messages: Array<Record<string, unknown>>): number {
  for (let i = messages.length - 1; i >= 0; i--) {
    const text = getMessageText(messages[i]);
    if (text.includes("[Pi Loop Node Context]")) return i;
  }
  return -1;
}

function getMessageText(message: Record<string, unknown>): string {
  const content = message.content;
  if (typeof content === "string") return content;
  if (Array.isArray(content)) {
    return content
      .map((part) => {
        if (part && typeof part === "object" && "text" in part) return String((part as { text?: unknown }).text ?? "");
        return "";
      })
      .join("\n");
  }
  return "";
}