import { readFileSync, existsSync, mkdirSync, writeFileSync, rmSync, openSync, fsyncSync, closeSync } from "node:fs";
import { join } from "node:path";
import { mkdtempSync } from "node:fs";
import os from "node:os";
import { Guardrails } from "./guardrails.js";

/**
 * Guardrails Replay CLI
 *
 * Replays a set of tool calls through the guardrails engine to calibrate rules.
 *
 * Usage:
 *   node guardrails-replay.mjs [--fixture FILE] [--url URL] [--directory DIR] [--limit N] [--assume-in-scope WT]
 */

/** Rule IDs that go in the "Always-Scope" table (always fire regardless of boulder). */
const ALWAYS_RULES = new Set(["bash_banned", "bash_protected_port"]);

/** Flush a file path to ensure disk sync. */
function _fsync(path) {
  try {
    const fd = openSync(path, "r");
    fsyncSync(fd);
    closeSync(fd);
  } catch { /* non-fatal */ }
}

/**
 * Extract worktree path from a prompt string.
 *
 * Parses the first `WORKTREE: <path>` line. Path is the first whitespace-delimited
 * token after `WORKTREE: `, with one trailing `.` stripped.
 * Returns null if no WORKTREE line is found.
 */
function parseWorktreePath(prompt) {
  if (typeof prompt !== "string") return null;
  const m = prompt.match(/WORKTREE:\s*(\S+)/m);
  if (!m) return null;
  let path = m[1];
  if (path.endsWith(".")) path = path.slice(0, -1);
  return path;
}

async function main() {
  const args = process.argv.slice(2);
  const options = {
    fixture: null,
    url: null,
    directory: process.cwd(),
    limit: null,
    assumeInScope: null,
    listAll: false,
  };

  for (let i = 0; i < args.length; i++) {
    if (args[i] === "--fixture") options.fixture = args[++i];
    else if (args[i] === "--url") options.url = args[++i];
    else if (args[i] === "--directory") options.directory = args[++i];
    else if (args[i] === "--limit") options.limit = parseInt(args[++i], 10);
    else if (args[i] === "--assume-in-scope") options.assumeInScope = args[++i];
    else if (args[i] === "--list-all") options.listAll = true;
  }

  const directory = options.directory;

  // Create a temp dir for synthetic config/boulder — never under --directory.
  const omoDir = mkdtempSync(join(os.tmpdir(), "guardrails-replay-"));

  // 1. Setup a "report" config (all rules in block mode, no logging)
  const reportConfig = {
    rules: {
      task_needs_agent: "block",
      task_banned_agent: "block",
      task_worktree_line: "block",
      write_outside_worktree: "block",
      bash_main_checkout: "block",
      bash_banned: "block",
      bash_protected_port: "block",
      plan_tick_gate: "block",
    },
    protected_ports: [8080],
  };
  writeFileSync(join(omoDir, "guardrails.json"), JSON.stringify(reportConfig));

  // 2. Load tool calls
  let calls = [];
  if (options.fixture) {
    const raw = readFileSync(options.fixture, "utf-8");
    calls = JSON.parse(raw);
  } else if (options.url) {
    const sessions = await fetchSessions(options.url, directory);
    for (const s of sessions) {
      const messages = await fetchMessages(options.url, s.id, directory);
      for (const m of messages) {
        if (m.parts) {
          for (const part of m.parts) {
            if (part.type === "tool" && part.state?.status === "completed") {
              calls.push({
                sessionID: s.id,
                callID: part.callID || part.partID || part.id || `part_${s.id}_${part.index}`,
                tool: part.tool,
                args: part.state?.input || {},
                prompt: part.state?.input?.prompt,
              });
            }
          }
        }
      }
    }
  }

  if (options.limit) calls = calls.slice(0, options.limit);

  // 3. Build per-session boulders.
  //    Group calls by sessionID. Each session gets its own boulder with
  //    worktree_path derived from its first WORKTREE: line (or --assume-in-scope).
  const sessionCalls = {};
  for (const call of calls) {
    const sid = call.sessionID;
    if (!sessionCalls[sid]) sessionCalls[sid] = [];
    sessionCalls[sid].push(call);
  }

  const sessionBoulders = {};
  for (const [sid, sessionCallsList] of Object.entries(sessionCalls)) {
    // Find first WORKTREE: line from any call in this session.
    let wtPath = null;
    for (const c of sessionCallsList) {
      wtPath = parseWorktreePath(c.prompt ?? "");
      if (wtPath) break;
    }

    // Fallback: --assume-in-scope.
    if (!wtPath) {
      wtPath = options.assumeInScope;
    }

    sessionBoulders[sid] = {
      status: "active",
      session_ids: [sid],
      worktree_path: wtPath || null,
    };
  }

  // Write all session boulders to omoDir *before* Guardrails reads them.
  for (const [sid, boulder] of Object.entries(sessionBoulders)) {
    const boulderPath = join(omoDir, `boulder-${sid}.json`);
    writeFileSync(boulderPath, JSON.stringify(boulder));
  }

  try {
    // 4. Replay — create a fresh Guardrails factory per session so readBoulder
    //    picks up that session's boulder from disk.  Group by err.ruleId:
    //      Always table = bash_banned + bash_protected_port
    //      (B) table    = the other six

    const stats = {
      always: {},
      scoped: {},
    };

    const blockedCalls = [];

    const client = {
      session: {
        get: async ({ path }) => ({ data: { id: path.id, parentID: null } }),
      },
    };

    for (const [sid, sessionCallList] of Object.entries(sessionCalls)) {
      const boulder = sessionBoulders[sid];

      // Write this session's boulder to omoDir *before* creating the factory.
      const boulderPath = join(omoDir, "boulder.json");
      writeFileSync(boulderPath, JSON.stringify(boulder));
      _fsync(boulderPath);

      // Create a fresh Guardrails factory so readBoulder picks up this boulder.
      const { "tool.execute.before": hook } = await Guardrails({
        client,
        directory,
        omoDir: omoDir,
      });

      for (const call of sessionCallList) {
        const input = { sessionID: sid, tool: call.tool, callID: call.callID };
        const output = { args: { ...call.args } };

        try {
          await hook(input, output);

          // Check if prompt was rewritten by task_worktree_line
          if (call.tool === "task" && output.args?.prompt !== call.args?.prompt) {
            updateStat("task_worktree_line", "rewrite", stats);
          }
        } catch (err) {
          const ruleId = err.ruleId;
          if (!ruleId) {
            console.error(`Unexpected error during replay of ${call.callID}:`, err);
            continue;
          }

          const type = "block";
          updateStat(ruleId, type, stats);

          if (options.listAll) {
            const content = call.args?.command || call.args?.description || "";
            const truncated = content.length > 160 ? content.slice(0, 160) : content;
            blockedCalls.push(`${ruleId} | ${call.tool} | ${truncated}`);
          }

          // Collect examples (max 5 per rule)
          const bucket = getBucket(ruleId, stats);
          if (!bucket.examples) bucket.examples = [];
          if (bucket.examples.length < 5) {
            bucket.examples.push(JSON.stringify(call.args));
          }
        }
      }
    }

    // 6. Print Results
    const printTable = (title, data) => {
      console.log(`\n${title}`);
      console.log("Rule | Blocked | Rewritten | Examples");
      console.log("-----|----------|-----------|----------");
      for (const [id, s] of Object.entries(data)) {
        console.log(`${id} | ${s.block} | ${s.rewrite || 0} | ${s.examples ? s.examples.join("; ") : ""}`);
      }
    };

    printTable("Always-Scope Rules", stats.always);
    printTable("(B) Scoped Rules", stats.scoped);

    if (options.listAll && blockedCalls.length > 0) {
      console.log("\nBlocked Calls Detail");
      console.log("Rule | Tool | Command/Description");
      console.log("----|------|---------------------");
      for (const line of blockedCalls) {
        console.log(line);
      }
    }
  } finally {
    rmSync(omoDir, { recursive: true, force: true });
  }
}

/** Get the stat bucket for a ruleId based on grouping. */
function getBucket(ruleId, stats) {
  const bucket = ALWAYS_RULES.has(ruleId)
    ? stats.always
    : stats.scoped;
  if (!bucket[ruleId]) {
    bucket[ruleId] = { block: 0, rewrite: 0, examples: [] };
  }
  return bucket[ruleId];
}

/** Increment a stat counter. */
function updateStat(ruleId, type, stats) {
  const bucket = getBucket(ruleId, stats);
  if (type === "block") bucket.block++;
  if (type === "rewrite") bucket.rewrite++;
}

/** Fetch sessions from the opencode API. */
async function fetchSessions(url, directory) {
  const params = new URLSearchParams({ directory });
  const resp = await fetch(`${url}/session?${params}`);
  if (!resp.ok) throw new Error(`GET ${url}/session failed: ${resp.status} ${resp.statusText}`);
  return resp.json();
}

/** Fetch messages for a session from the opencode API. */
async function fetchMessages(url, sessionId, directory) {
  const params = new URLSearchParams({ directory });
  const resp = await fetch(`${url}/session/${sessionId}/message?${params}`);
  if (!resp.ok) throw new Error(`GET ${url}/session/${sessionId}/message failed: ${resp.status} ${resp.statusText}`);
  return resp.json();
}

main().catch(console.error);