import { stat } from "node:fs/promises";
import {
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
  type ExtensionAPI,
  formatSize,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { configLoader } from "../../src/shared/config";
import {
  createFeatureRegisterPayload,
  emitConfigReload,
  emitWorkspaceRootsChanged,
  GUARDRAILS_FEATURE_REGISTER_EVENT,
  GUARDRAILS_FEATURE_REQUEST_EVENT,
  GUARDRAILS_YOLO_CHANGED_EVENT,
  type GuardrailsYoloChangedPayload,
} from "../../src/shared/events";
import {
  createAddDirAutocompleteProvider,
  createRootExpansionProvider,
} from "./autocomplete";
import { createNoAccessFilter, filterRipgrepOutput } from "./policy-filter";
import {
  describeEntry,
  listRemovableEntries,
  removeConfigRoot,
  removePathGrant,
} from "./remove";
import {
  formatWorkspaceRoots,
  persistRoot,
  resolveRoots,
  resolveTilde,
  type WorkspaceRoot,
} from "./roots";
import {
  buildRipgrepArgs,
  resolveRipgrepOutput,
  truncateWorkspaceOutput,
} from "./search";

export default async function workspaceRoots(pi: ExtensionAPI) {
  await configLoader.load();

  let cwd = process.cwd();
  let sessionRoots: WorkspaceRoot[] = [];
  let yoloEnabled = false;

  pi.events.on(GUARDRAILS_YOLO_CHANGED_EVENT, (data: unknown) => {
    const payload = data as GuardrailsYoloChangedPayload | undefined;
    if (typeof payload?.enabled === "boolean") yoloEnabled = payload.enabled;
  });

  const featureEnabled = () => {
    const config = configLoader.getConfig();
    return config.enabled && config.features.workspaceRoots;
  };

  /** Configured roots (all scopes) plus session-only /add-dir roots. */
  const currentRoots = (): WorkspaceRoot[] => {
    if (!featureEnabled()) return [];
    const configured = resolveRoots(
      configLoader.getConfig().workspaceRoots.roots,
      cwd,
    );
    const known = new Set(configured.map((r) => r.path));
    return [...configured, ...sessionRoots.filter((r) => !known.has(r.path))];
  };

  const announce = () => {
    emitWorkspaceRootsChanged(
      pi,
      currentRoots().map((r) => r.path),
    );
  };

  pi.events.on(GUARDRAILS_FEATURE_REQUEST_EVENT, () => {
    pi.events.emit(
      GUARDRAILS_FEATURE_REGISTER_EVENT,
      createFeatureRegisterPayload("workspaceRoots"),
    );
  });

  const textResult = (text: string, details: Record<string, unknown> = {}) => ({
    content: [{ type: "text" as const, text }],
    details,
  });

  pi.on("session_start", async (_event, ctx) => {
    cwd = ctx.cwd;
    announce();

    ctx.ui.addAutocompleteProvider((current) =>
      createAddDirAutocompleteProvider(current, ctx.cwd),
    );
    ctx.ui.addAutocompleteProvider((current) =>
      createRootExpansionProvider(current, currentRoots),
    );
  });

  pi.on("before_agent_start", (event) => {
    const roots = currentRoots();
    if (roots.length === 0) return {};
    return { systemPrompt: event.systemPrompt + formatWorkspaceRoots(roots) };
  });

  pi.registerCommand("add-dir", {
    description:
      "Add an external directory as a workspace root (known to the agent and allowed by path access), or list current roots",
    handler: async (args, ctx) => {
      if (!featureEnabled()) {
        ctx.ui.notify(
          "Workspace roots are disabled (features.workspaceRoots).",
          "warning",
        );
        return;
      }

      const rawArg = args.trim();
      if (!rawArg) {
        const roots = currentRoots();
        if (roots.length === 0) {
          ctx.ui.notify(
            "No workspace roots configured.\nType /add-dir  (with a space) and use autocomplete to browse for a directory.",
            "info",
          );
          return;
        }
        const lines = roots.map(
          (r) => `  ${r.path}${r.alias ? `  (${r.alias})` : ""}`,
        );
        ctx.ui.notify(`Workspace roots:\n${lines.join("\n")}`, "info");
        return;
      }

      const path = resolveTilde(rawArg, cwd);
      const stats = await stat(path).catch(() => null);
      if (!stats?.isDirectory()) {
        ctx.ui.notify(`Not a directory: ${path}`, "error");
        return;
      }
      if (currentRoots().some((r) => r.path === path)) {
        ctx.ui.notify(`Already a workspace root: ${path}`, "info");
        return;
      }

      const SESSION = "Session only";
      const ALWAYS = "Always (save to project guardrails config)";
      const choice = ctx.hasUI
        ? await ctx.ui.select(`Add workspace root: ${path}`, [SESSION, ALWAYS])
        : SESSION;
      if (choice === undefined) {
        ctx.ui.notify(`Cancelled — root not added: ${path}`, "info");
        return;
      }

      if (choice === ALWAYS) {
        await persistRoot({ path });
        ctx.ui.notify(
          `Added workspace root (saved to project config): ${path}`,
          "info",
        );
      } else {
        sessionRoots = [...sessionRoots, { path }];
        ctx.ui.notify(`Added workspace root for this session: ${path}`, "info");
      }
      announce();
    },
  });

  pi.registerCommand("remove-dir", {
    description:
      "Remove a workspace root or persisted path-access directory grant, or pick one from a list",
    handler: async (args, ctx) => {
      if (!featureEnabled()) {
        ctx.ui.notify(
          "Workspace roots are disabled (features.workspaceRoots).",
          "warning",
        );
        return;
      }

      // Reload from disk so grants path-access persisted mid-session show up.
      await configLoader.load();
      let entries = listRemovableEntries(sessionRoots);
      if (entries.length === 0) {
        ctx.ui.notify(
          "No workspace roots or directory grants to remove.",
          "info",
        );
        return;
      }

      const rawArg = args.trim();
      if (rawArg) {
        const resolved = resolveTilde(rawArg, cwd);
        entries = entries.filter(
          (entry) =>
            entry.path === rawArg || resolveTilde(entry.path, cwd) === resolved,
        );
        if (entries.length === 0) {
          ctx.ui.notify(`No allowed directory matches: ${rawArg}`, "error");
          return;
        }
      }

      let chosen = rawArg && entries.length === 1 ? entries[0] : undefined;
      if (!chosen) {
        if (!ctx.hasUI) {
          ctx.ui.notify(
            `Specify which directory to remove:\n${entries
              .map((entry) => `  ${describeEntry(entry)}`)
              .join("\n")}`,
            "info",
          );
          return;
        }
        const labels = entries.map(describeEntry);
        const choice = await ctx.ui.select(
          "Remove which allowed directory?",
          labels,
        );
        if (choice === undefined) return;
        chosen = entries[labels.indexOf(choice)];
        if (!chosen) return;
      }

      if (chosen.kind === "session-root") {
        const { path } = chosen;
        sessionRoots = sessionRoots.filter((root) => root.path !== path);
      } else if (chosen.kind === "config-root") {
        await removeConfigRoot(chosen.scope, chosen.path);
      } else {
        await removePathGrant(chosen.scope, chosen.path);
        emitConfigReload(pi);
      }
      announce();
      ctx.ui.notify(`Removed — ${describeEntry(chosen)}`, "info");
    },
  });

  pi.registerTool({
    name: "search_workspaces",
    label: "Search Workspaces",
    description: `Search for a pattern across all configured workspace roots using ripgrep. Returns matching file paths with context lines. Output is limited to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)}; complete truncated output is saved to a temporary file. Use this when you need to find files or code references that might be in any workspace root.`,
    promptSnippet:
      "Search across all workspace roots for files matching a pattern",
    promptGuidelines: [
      "Use search_workspaces when you suspect relevant code or files might be in another workspace root.",
      "search_workspaces searches globs, file names, and file contents using ripgrep.",
    ],
    parameters: Type.Object({
      pattern: Type.String({
        description: "Pattern to search for (regular expression)",
      }),
      glob: Type.Optional(
        Type.String({
          description:
            "Glob pattern to filter files (e.g. '*.ts', 'src/**/*.rs')",
        }),
      ),
    }),
    async execute(_toolCallId, params, signal, onUpdate) {
      const roots = currentRoots().map((r) => r.path);
      if (roots.length === 0) {
        return textResult(
          "No workspace roots configured. Add them with /add-dir or in guardrails.json under workspaceRoots.roots.",
        );
      }
      onUpdate?.(textResult("Searching workspace roots..."));

      const result = await pi.exec(
        "rg",
        buildRipgrepArgs(params.pattern, roots, params.glob),
        { signal, timeout: 15000 },
      );
      const filtered = filterRipgrepOutput(
        resolveRipgrepOutput(result),
        yoloEnabled ? () => true : createNoAccessFilter(cwd),
      );
      const notice = filtered.hidden
        ? `\n\n[${filtered.hidden} match(es) in files protected by guardrails policies were hidden.]`
        : "";
      const output = await truncateWorkspaceOutput(filtered.text + notice);

      return textResult(output.text, {
        roots,
        pattern: params.pattern,
        glob: params.glob,
        exitCode: result.code,
        hiddenByPolicy: filtered.hidden,
        truncated: output.truncated,
        fullOutputPath: output.fullOutputPath,
      });
    },
  });

  pi.registerTool({
    name: "search_workspace_files",
    label: "Search Workspace Files",
    description:
      "Find files by name across all configured workspace roots. Use this to locate a specific file when you know its name but not which workspace root it lives in.",
    promptSnippet: "Find files by name across all workspace roots",
    promptGuidelines: [
      "Use search_workspace_files when the user references a file by name and you do not know which workspace root it is in.",
    ],
    parameters: Type.Object({
      name: Type.String({
        description:
          "File name or glob to search for (e.g. 'auth.ts', '*.schema.json')",
      }),
    }),
    async execute(_toolCallId, params, signal, onUpdate) {
      const roots = currentRoots();
      if (roots.length === 0)
        return textResult("No workspace roots configured.");
      onUpdate?.(textResult("Finding files across roots..."));

      const allowPath = yoloEnabled ? () => true : createNoAccessFilter(cwd);
      // fd ships with pi and works on Windows, where PATH lookup would find
      // System32's find.exe (a text matcher) instead of a file finder.
      // --hidden --no-ignore preserves the old `find -name` semantics; output
      // paths are absolute because root.path is absolute.
      const matches = (
        await Promise.all(
          roots.map((root) =>
            pi
              .exec(
                "fd",
                [
                  "--type",
                  "f",
                  "--glob",
                  "--hidden",
                  "--no-ignore",
                  params.name,
                  root.path,
                ],
                { signal, timeout: 10000 },
              )
              .then((r) => r.stdout.split("\n").filter(Boolean))
              .catch(() => []),
          ),
        )
      )
        .flat()
        .filter(allowPath);

      const output = await truncateWorkspaceOutput(
        matches.length
          ? matches.join("\n")
          : `No files matching '${params.name}' found in workspace roots.`,
      );
      return textResult(output.text, {
        name: params.name,
        matches: matches.length,
        truncated: output.truncated,
        fullOutputPath: output.fullOutputPath,
      });
    },
  });
}