import {
  type ExtensionAPI,
  type ExtensionContext,
  isToolCallEventType
} from "@earendil-works/pi-coding-agent";
import * as path from "node:path";
import * as fs from "node:fs";
import { execSync } from "node:child_process";

// ═══════════════════════════════════════════════════════════════════════════
// guardrails.ts — Single interception pipeline that runs before every tool
// call and decides: allow it, block it, or require confirmation.
//
// Pipeline order (A → B → C → D):
//   A) checkScope        — workspace boundary enforcement
//   B) checkProtectedPath — .env, .git, node_modules, build output, lockfiles
//   C) checkRiskyCommand  — dangerous bash patterns (force-push, rm -rf, etc.)
//   D) checkDiffPreview   — diff-style approval for file-modifying operations
//
// Order rationale: scope is cheapest and most important (no point previewing
// a diff for a path outside the workspace). Protected-path check is the next
// cheapest hard check. Risky command patterns are cheap too but need dirtyTree
// from session start. Diff preview is the most expensive (reads files) and
// should only run after all hard/soft permission checks pass.
// ═══════════════════════════════════════════════════════════════════════════

// ─── Configuration ───────────────────────────────────────────────────────

/** Protected path rules. `blockRead`/`blockWrite` control whether the action
 *  is blocked (true), requires CONFIRM ("confirm"), or allowed (false). */
const PROTECTED_PATH_RULES: Array<{
  pattern: RegExp;
  name: string;
  blockRead: boolean;
  blockWrite: boolean | "confirm";
}> = [
  { pattern: /(^|[/\\])\.env([/\\]|$)/i,         name: ".env file",              blockRead: true,  blockWrite: true },
  { pattern: /(^|[/\\])\.env\.\w+$/i,               name: ".env file",            blockRead: true,  blockWrite: true },
  { pattern: /(^|[/\\])\.git[/\\]/i,                name: ".git internals",        blockRead: false, blockWrite: true },
  { pattern: /(^|[/\\])node_modules[/\\]/i,          name: "node_modules",          blockRead: false, blockWrite: true },
  { pattern: /(^|[/\\])(dist|build|out|coverage)[/\\]/i, name: "build output",     blockRead: false, blockWrite: true },
  { pattern: /(^|[/\\])package-lock\.json$/i,       name: "lockfile",              blockRead: false, blockWrite: "confirm" },
  { pattern: /(^|[/\\])yarn\.lock$/i,               name: "lockfile",              blockRead: false, blockWrite: "confirm" },
  { pattern: /(^|[/\\])pnpm-lock\.yaml$/i,          name: "lockfile",              blockRead: false, blockWrite: "confirm" },
  { pattern: /(^|[/\\])bun\.lockb$/i,               name: "lockfile",              blockRead: false, blockWrite: "confirm" },
];

/** Risky bash command patterns. `destructive` means uncommitted work may be lost. */
const RISKY_COMMAND_PATTERNS: Array<{
  pattern: RegExp;
  name: string;
  destructive: boolean;
}> = [
  { pattern: /\bgit\s+push\b.*?(?:-f|--force)/i,         name: "git force-push",                  destructive: true },
  { pattern: /\bgit\s+reset\s+--hard\b/i,                  name: "git reset --hard",                destructive: true },
  { pattern: /\bgit\s+clean\s+-fd/i,                        name: "git clean -fd",                    destructive: true },
  { pattern: /\brm\s+-rf\b/i,                               name: "recursive force delete",          destructive: false },
  { pattern: /\bchmod\s+[0-7]{3,4}\s+/i,                   name: "permission change",               destructive: false },
  { pattern: /\bchown\b/i,                                  name: "ownership change",                destructive: false },
  { pattern: /\bmkfs\b|\bdd\s+if=/i,                        name: "filesystem/destructive disk op", destructive: true },
  { pattern: /\bsudo\b/i,                                   name: "sudo command",                    destructive: false },
  { pattern: /\bcurl\b.*?\|\s*bash\b/i,                    name: "pipe-to-shell",                   destructive: false },
  { pattern: /\bwget\b.*?-O\s*-\s*\|\s*bash\b/i,          name: "pipe-to-shell",                   destructive: false },
  { pattern: /\bdangerous-delete|force-delete\b/i,          name: "dangerous delete",                destructive: true },
];

/** Patterns that indicate a bash command may modify the filesystem. */
const FILE_MODIFYING_BASH_PATTERNS: RegExp[] = [
  />{1,2}\s+/,
  /\btee\b/,
  /\bsed\s+-i\b/,
  /\bmv\b/,
  /\brm\b/,
  /\bcp\b/,
  /\bmkdir\b/,
  /\btouch\b/,
  /\bchmod\b/,
  /\bchown\b/,
  /\bgit\s+(add|commit|checkout|reset|apply|merge|rebase|push)\b/,
  /\bnpm\s+(install|uninstall|run)\b/,
  /\bpip\s+install\b/,
];

// ─── Module-level state (immutable after session_start) ─────────────────

let workspaceRoot: string | null = null;
let dirtyTree: boolean = false;
let suppressConfirmations: boolean = false;

// ─── GuardResult type ────────────────────────────────────────────────────

/** Discriminated union for all guard-check results.
 *  - `{ allow: true }` → proceed
 *  - `{ allow: false; block: true; reason }` → hard block, no override
 *  - `{ allow: false; confirm: true; prompt }` → require user confirmation */
type GuardResult =
  | { allow: true }
  | { allow: false; block: true; reason: string }
  | { allow: false; confirm: true; prompt: string };

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

/** Resolve a possibly-relative path against the locked workspace root. */
function resolvePath(p: string): string {
  if (!workspaceRoot) return path.resolve(p);
  if (path.isAbsolute(p)) return path.normalize(p);
  return path.resolve(workspaceRoot, p);
}

/** True if targetPath lies inside the locked workspace root. */
function isInsideWorkspace(targetPath: string): boolean {
  if (!workspaceRoot) return true;
  const resolved = resolvePath(targetPath);
  const relative = path.relative(workspaceRoot, resolved);
  // relative is empty for exact match, starts with '' for subdirs
  return relative === "" || (!relative.startsWith("..") && !path.isAbsolute(relative));
}

/** Quick git-dirty check via `git status --porcelain`. Returns true if clean. */
function isGitClean(dir: string): boolean {
  try {
    const out = execSync("git status --porcelain", {
      cwd: dir,
      encoding: "utf-8",
      stdio: ["ignore", "pipe", "pipe"],
      timeout: 5000,
    });
    return out.trim().length === 0;
  } catch {
    return true; // not a git repo, git unavailable, or timeout — assume clean
  }
}

/** Extract file paths from a tool call's input, depending on tool type. */
function getToolPaths(toolName: string, input: Record<string, unknown>): string[] {
  switch (toolName) {
    case "read":
    case "write":
      return [input.path as string].filter(Boolean);
    case "edit":
      return [input.filePath as string].filter(Boolean);
    case "multi_edit": {
      const edits = input.edits as Array<{ filePath: string }> | undefined;
      return (edits ?? []).map((e) => e.filePath).filter(Boolean);
    }
    case "ls":
    case "find":
    case "grep":
      return [input.path as string].filter(Boolean);
    default:
      return [];
  }
}

/** True if the tool call could modify the filesystem. */
function isWriteOperation(toolName: string, input: Record<string, unknown>): boolean {
  if (["write", "edit", "multi_edit"].includes(toolName)) return true;
  if (toolName === "bash") {
    return isFileModifyingBash(input.command as string ?? "");
  }
  return false;
}

/** True if a bash command looks like it modifies files. */
function isFileModifyingBash(command: string): boolean {
  return FILE_MODIFYING_BASH_PATTERNS.some((p) => p.test(command));
}

// ─── Pipeline Step A — Scope Check ─────────────────────────────────────
/**
 * Runs FIRST, before any other check. Determines whether the target of the
 * tool call lies inside the locked workspace root. If not, BLOCK unconditionally
 * — no confirmation path. This is the cheapest and most critical gate.
 *
 * Why first: there is no point checking protection rules or generating diffs
 * for a path outside the workspace; the answer is always "no", and it is the
 * fastest check.
 */
function checkScope(
  toolName: string,
  input: Record<string, unknown>,
): { allow: true } | { allow: false; block: true; reason: string } {
  if (!workspaceRoot) return { allow: true };

  // For bash commands, check for explicit directory changes outside workspace
  if (toolName === "bash") {
    const command = (input.command as string) ?? "";

    // Check if command has a `cd` to an absolute path or `..` escape
    // that resolves outside the workspace.
    const cdMatch = command.match(/\bcd\s+(\S+)/);
    if (cdMatch) {
      const cdTarget = cdMatch[1];
      // If it's an absolute path or a relative path with `..`, resolve it
      if (cdTarget.startsWith("/") || cdTarget.startsWith("..")) {
        const resolved = resolvePath(cdTarget);
        if (!isInsideWorkspace(cdTarget)) {
          return {
            allow: false,
            block: true,
            reason: `Scope blocked: \`cd\` target "${cdTarget}" resolves outside workspace root "${workspaceRoot}". All operations must stay within the project directory.`,
          };
        }
      }
    }

    // Check for absolute paths in arguments to common file commands
    // that could target outside workspace
    const fileArgMatch = command.match(
      /\b(?:cat|less|more|head|tail|grep|find|ls|rm|cp|mv|chmod|chown|touch|mkdir)\s+(\/\S+)/,
    );
    if (fileArgMatch) {
      const absPath = fileArgMatch[1];
      if (!isInsideWorkspace(absPath)) {
        return {
          allow: false,
          block: true,
          reason: `Scope blocked: argument "${absPath}" is an absolute path outside workspace root "${workspaceRoot}". All operations must stay within the project directory.`,
        };
      }
    }

    return { allow: true };
  }

  // For non-bash tools, check each path from the input
  const paths = getToolPaths(toolName, input);
  for (const p of paths) {
    if (!isInsideWorkspace(p)) {
      return {
        allow: false,
        block: true,
        reason: `Scope blocked: path "${p}" resolves outside workspace root "${workspaceRoot}". All operations must stay within the project directory.`,
      };
    }
  }

  return { allow: true };
}

// ─── Pipeline Step B — Protected Path Check ────────────────────────────
/**
 * Runs AFTER scope check (A). Checks whether the target path matches any
 * protected path rule (.env, .git internals, node_modules, build output,
 * lockfiles).
 *
 * Why after A: paths that fail scope check don't need protection rules.
 * Why before C: path matching is cheaper than regex on arbitrary bash commands.
 *
 * - .env (read OR write) → BLOCK, no confirmation, no override
 * - .git, node_modules, build output (write) → BLOCK, no confirmation
 * - lockfiles (write) → CONFIRM (user may approve)
 */
function checkProtectedPath(
  toolName: string,
  input: Record<string, unknown>,
): GuardResult {
  const isWrite = isWriteOperation(toolName, input);
  const paths = getToolPaths(toolName, input);

  for (const p of paths) {
    if (!p) continue;
    // Normalise separators for consistent pattern matching
    const normalized = p.replace(/\\/g, "/");

    for (const rule of PROTECTED_PATH_RULES) {
      if (!rule.pattern.test(normalized)) continue;

      // Read action
      if (!isWrite) {
        if (rule.blockRead) {
          return { allow: false, block: true, reason: `Blocked: ${rule.name} ("${p}") — read is not allowed.` };
        }
        continue; // read of this path is allowed
      }

      // Write action
      if (rule.blockWrite === true) {
        return { allow: false, block: true, reason: `Blocked: ${rule.name} ("${p}") — editing is not allowed.` };
      }
      if (rule.blockWrite === "confirm") {
        return {
          allow: false,
          confirm: true,
          prompt: `Edit ${rule.name} "${p}"?\n\nModifying lockfiles can cause dependency conflicts. Confirm only if you intend to change dependency versions or update the lockfile intentionally.`,
        };
      }
    }
  }

  return { allow: true };
}

// ─── Pipeline Step C — Risky Shell Command Check ──────────────────────
/**
 * Runs AFTER scope check (A) and protected-path check (B). Checks bash
 * commands against risky patterns (force-push, reset --hard, rm -rf, etc.).
 *
 * Why after B: the protected-path check catches .env and similar rules
 * that apply to all tools. This check is bash-only.
 *
 * If the command matches a destructive pattern AND the working tree is dirty,
 * the confirmation prompt explicitly warns about uncommitted changes.
 */
function checkRiskyCommand(
  toolName: string,
  input: Record<string, unknown>,
): { allow: true } | { allow: false; confirm: true; prompt: string } {
  if (toolName !== "bash") return { allow: true };

  const command = (input.command as string) ?? "";

  for (const rule of RISKY_COMMAND_PATTERNS) {
    if (rule.pattern.test(command)) {
      let prompt = `Run risky command: ${rule.name}?`;
      if (rule.destructive && dirtyTree) {
        prompt += `\n\n⚠️  WARNING: working tree has uncommitted changes that may be lost.`;
      }
      return { allow: false, confirm: true, prompt };
    }
  }

  return { allow: true };
}

// ─── Pipeline Step D — Diff Preview for File-Modifying Ops ────────────
/**
 * Runs LAST, after all prior checks pass (A–C). Generates a before/after
 * diff-style preview for write, edit, multi_edit, and file-modifying bash
 * commands. Requires explicit user approval.
 *
 * Why last: this is the most expensive step (reads current file content,
 * computes diffs). It is the final "do you want this specific change" gate.
 *
 * If `suppressConfirmations` is true, this step is skipped.
 */
async function checkDiffPreview(
  toolName: string,
  input: Record<string, unknown>,
  ctx: ExtensionContext,
): Promise<{ allow: true } | { allow: false; confirm: true; prompt: string }> {
  if (suppressConfirmations) return { allow: true };

  // ── write tool ──────────────────────────────────────────────────────
  if (toolName === "write") {
    return previewWrite(input as { path: string; content: string });
  }

  // ── edit tool ───────────────────────────────────────────────────────
  if (toolName === "edit") {
    return previewEdit(input as { filePath: string; oldText: string; newText: string });
  }

  // ── multi_edit tool ─────────────────────────────────────────────────
  if (toolName === "multi_edit") {
    return previewMultiEdit(input as { edits: Array<{ filePath: string; oldText: string; newText: string }> });
  }

  // ── bash (file-modifying) ───────────────────────────────────────────
  if (toolName === "bash") {
    return previewBash(input as { command: string });
  }

  return { allow: true };
}

/** Generate a diff preview for the `write` tool. */
function previewWrite(input: { path: string; content: string }): { allow: false; confirm: true; prompt: string } {
  const filePath = input.path;
  const newContent = input.content;

  let existing: string | null = null;
  try {
    existing = fs.readFileSync(resolvePath(filePath), "utf-8");
  } catch {
    existing = null; // new file
  }

  if (existing !== null) {
    return {
      allow: false,
      confirm: true,
      prompt: formatDiff(filePath, existing, newContent),
    };
  }

  // New file: show preview of content
  const lines = newContent.split("\n");
  const previewLines = lines.slice(0, 30);
  let text = `Create new file: ${filePath}\n`;
  text += previewLines.map((l) => `+ ${l}`).join("\n");
  if (lines.length > 30) {
    text += `\n... (${lines.length - 30} more lines)`;
  }
  return { allow: false, confirm: true, prompt: text };
}

/** Generate a diff preview for the `edit` tool. */
function previewEdit(input: { filePath: string; oldText: string; newText: string }): { allow: false; confirm: true; prompt: string } {
  const filePath = input.filePath;
  const oldText = input.oldText;
  const newText = input.newText;

  let existing: string | null = null;
  try {
    existing = fs.readFileSync(resolvePath(filePath), "utf-8");
  } catch {
    existing = null;
  }

  if (existing === null) {
    // File doesn't exist — show what would be inserted
    return {
      allow: false,
      confirm: true,
      prompt: `Edit ${filePath} (file does not exist yet):\n+ ${newText.split("\n").join("\n+ ")}`,
    };
  }

  if (!existing.includes(oldText)) {
    const snippet = existing.slice(0, 300);
    return {
      allow: false,
      confirm: true,
      prompt: `Edit ${filePath}: ⚠️  oldText not found in current file. Showing beginning of file:\n${snippet}${existing.length > 300 ? "\n..." : ""}`,
    };
  }

  return {
    allow: false,
    confirm: true,
    prompt: formatDiff(filePath, oldText, newText),
  };
}

/** Generate a combined preview for `multi_edit`. */
function previewMultiEdit(
  input: { edits: Array<{ filePath: string; oldText: string; newText: string }> },
): { allow: true } | { allow: false; confirm: true; prompt: string } {
  const edits = input.edits ?? [];
  if (edits.length === 0) return { allow: true };

  const parts: string[] = [`Multi-edit preview (${edits.length} file${edits.length > 1 ? "s" : ""}):`];
  for (const edit of edits) {
    const oldFirst = edit.oldText.split("\n")[0].slice(0, 80);
    const newFirst = edit.newText.split("\n")[0].slice(0, 80);
    parts.push(`\n  ${edit.filePath}:`);
    if (oldFirst) parts.push(`    - ${oldFirst}`);
    if (newFirst) parts.push(`    + ${newFirst}`);
  }

  return { allow: false, confirm: true, prompt: parts.join("\n") };
}

/** Show the command itself as a preview for file-modifying bash commands. */
function previewBash(input: { command: string }): { allow: false; confirm: true; prompt: string } {
  const command = input.command;
  const truncated = command.length > 500;
  return {
    allow: false,
    confirm: true,
    prompt: `Run file-modifying command:\n\n$ ${command.slice(0, 500)}${truncated ? "\n... (truncated)" : ""}`,
  };
}

/** Produce a simple unified-diff-style string for a pair of texts. */
function formatDiff(label: string, oldText: string, newText: string): string {
  const oldLines = oldText.split("\n");
  const newLines = newText.split("\n");
  const lines: string[] = [`Diff for ${label}:`];

  // Simple line-by-line diff
  let inDiff = false;
  const maxLen = Math.max(oldLines.length, newLines.length);
  for (let i = 0; i < maxLen; i++) {
    const oldLine = i < oldLines.length ? oldLines[i] : undefined;
    const newLine = i < newLines.length ? newLines[i] : undefined;
    if (oldLine === newLine) {
      if (inDiff) {
        // Separator line after diff block
        // (no separator needed; context lines indicate boundary)
      }
      continue;
    }
    if (oldLine !== undefined) lines.push(`  - ${oldLine}`);
    if (newLine !== undefined) lines.push(`  + ${newLine}`);
    inDiff = true;
  }

  if (lines.length <= 1) {
    lines.push("  (no differences)");
  }

  return lines.join("\n");
}

// ─── User confirmation helper ────────────────────────────────────────────

async function confirmWithUser(
  ctx: ExtensionContext,
  prompt: string,
): Promise<boolean> {
  if (ctx.hasUI) {
    return !!(await ctx.ui.confirm("Guardrails", prompt));
  }
  // No UI available — block by default for safety
  return false;
}

// ═══════════════════════════════════════════════════════════════════════════
// Extension Entry Point
// ═══════════════════════════════════════════════════════════════════════════

export default function guardrailsExtension(pi: ExtensionAPI): void {
  // ── Session Start ──────────────────────────────────────────────────
  // Lock the workspace root and snapshot git dirty state.
  // Show a one-time non-blocking warning if the tree is dirty.
  pi.on("session_start", async (_event, ctx) => {
    workspaceRoot = ctx.cwd;
    dirtyTree = !isGitClean(workspaceRoot);

    if (dirtyTree && ctx.hasUI) {
      ctx.ui.notify(
        "Guardrails: session started with uncommitted changes. Some risky-git confirmations will mention lost work risk.",
        "warning",
      );
    }
  });

  // ── Tool Call Interception Pipeline ─────────────────────────────────
  // Every tool call passes through A → B → C → D in order.
  // Stop on first BLOCK. Pause for user on first CONFIRM.
  // Only proceed to the next check (or to execution) when current
  // check resolves to ALLOW.
  pi.on("tool_call", async (event, ctx) => {
    const { toolName, input } = event;

    // ── Step A: Scope check ────────────────────────────────────────────
    const scopeResult = checkScope(toolName, input);
    if (!scopeResult.allow) {
      return { block: true, reason: scopeResult.reason };
    }

    // ── Step B: Protected-path check ───────────────────────────────────
    const protectedResult = checkProtectedPath(toolName, input);
    if (!protectedResult.allow) {
      if ("block" in protectedResult && protectedResult.block) {
        return { block: true, reason: protectedResult.reason };
      }
      // CONFIRM branch
      if ("confirm" in protectedResult && protectedResult.confirm) {
        const approved = await confirmWithUser(ctx, protectedResult.prompt);
        if (!approved) {
          return { block: true, reason: "Blocked by user: protected path edit not approved." };
        }
      }
    }

    // ── Step C: Risky shell command check ──────────────────────────────
    const riskyResult = checkRiskyCommand(toolName, input);
    if (!riskyResult.allow && "confirm" in riskyResult && riskyResult.confirm) {
      const approved = await confirmWithUser(ctx, riskyResult.prompt);
      if (!approved) {
        return { block: true, reason: "Blocked by user: risky command not approved." };
      }
    }

    // ── Step D: Diff preview ───────────────────────────────────────────
    const diffResult = await checkDiffPreview(toolName, input, ctx);
    if (!diffResult.allow && "confirm" in diffResult && diffResult.confirm) {
      const approved = await confirmWithUser(ctx, diffResult.prompt);
      if (!approved) {
        return { block: true, reason: "Blocked by user: change not approved." };
      }
    }

    // ── Post-execution bookkeeping ─────────────────────────────────────
    // Re-check dirty state after any tree-modifying operation
    if (isWriteOperation(toolName, input) && workspaceRoot) {
      dirtyTree = !isGitClean(workspaceRoot);
    }

    // Allow the tool to proceed (return undefined = no block)
    return;
  });
}