import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

const PATH_TOOLS = new Set(["read", "write", "edit", "grep", "find", "ls"]);

function isWithin(root: string, candidate: string): boolean {
  const relative = path.relative(root, candidate);
  return relative === "" || (!relative.startsWith(`..${path.sep}`) && relative !== ".." && !path.isAbsolute(relative));
}

function canonicalizePotentialPath(inputPath: string, cwd: string): string {
  let normalized = inputPath.startsWith("@") ? inputPath.slice(1) : inputPath;
  if (normalized === "~") normalized = os.homedir();
  else if (normalized.startsWith(`~${path.sep}`)) normalized = path.join(os.homedir(), normalized.slice(2));

  const resolved = path.resolve(cwd, normalized);
  const missingSegments: string[] = [];
  let existing = resolved;

  while (!fs.existsSync(existing)) {
    const parent = path.dirname(existing);
    if (parent === existing) break;
    missingSegments.unshift(path.basename(existing));
    existing = parent;
  }

  const canonicalBase = fs.existsSync(existing) ? fs.realpathSync(existing) : existing;
  return path.join(canonicalBase, ...missingSegments);
}

export function registerWorkerGuard(pi: ExtensionAPI, configuredRoot: string): void {
  const root = fs.realpathSync(configuredRoot);

  pi.on("tool_call", async (event) => {
    if (!PATH_TOOLS.has(event.toolName)) return;

    const input = event.input as { path?: unknown };
    if (input.path === undefined && (event.toolName === "grep" || event.toolName === "find" || event.toolName === "ls")) {
      return;
    }
    if (typeof input.path !== "string") {
      return { block: true, reason: `${event.toolName} requires a path inside the assigned worktree.` };
    }

    const target = canonicalizePotentialPath(input.path, root);
    if (!isWithin(root, target)) {
      return {
        block: true,
        reason: `Worker file tools are restricted to ${root}. Attempted path: ${input.path}`,
      };
    }
  });

  pi.on("before_agent_start", async (event) => ({
    systemPrompt:
      event.systemPrompt +
      `\n\nYou are running as a worktree worker. Your assigned repository root is ${root}. ` +
      "Keep all file mutations and shell work inside this worktree. Do not access or modify sibling worktrees.",
  }));
}