// LSP Extension - Registers tools that let the LLM query language servers
// for hover, definition, references, completions, document symbols, and
// diagnostics. Tool calls are forwarded to the long-lived pi-lsp daemon
// (autospawned on first use) so LSP servers stay warm across calls.
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import * as path from "node:path";
import { uriToPath } from "./src/client.ts";
import {
  daemonDestroyServer,
  daemonDiagnostics,
  daemonRequest,
  daemonStatus,
} from "./src/daemonClient.ts";
import { pickServer, isServerAvailable } from "./src/root.ts";
import { servers } from "./server.ts";
import {
  ServerNotFoundError,
  UnsupportedExtensionError,
} from "./src/types.ts";

// Format Hover - Turn an LSP hover response into readable text.
function formatHover(result: unknown): string {
  if (!result || typeof result !== "object") return "(no hover info)";
  const hover = result as { contents?: unknown };
  if (!hover.contents) return "(empty)";

  // MarkupContent
  const contents = hover.contents as Record<string, unknown>;
  if (
    "value" in contents &&
    typeof contents.value === "string"
  ) {
    return contents.value;
  }
  // MarkedString | MarkedString[]
  if (Array.isArray(hover.contents)) {
    return hover.contents
      .map((s: any) => (typeof s === "string" ? s : (s?.value ?? "")))
      .join("\n");
  }
  if (
    "value" in contents &&
    typeof contents.language === "string"
  ) {
    return `\`\`\`${contents.language}\n${contents.value}\n\`\`\``;
  }
  return JSON.stringify(result, null, 2);
}

// Format Definition - Turn definition locations into readable text.
function formatDefinition(result: unknown): string {
  if (!result) return "(no definition found)";
  const locations = Array.isArray(result) ? result : [result];
  if (locations.length === 0) return "(no definition found)";

  return locations
    .map((loc: any, i: number) => {
      const file = uriToPath(loc.uri);
      const range = loc.range;
      return `${i + 1}. ${file} (${range.start.line + 1}:${range.start.character + 1})`;
    })
    .join("\n");
}

// Format References - Turn reference locations into readable text.
function formatReferences(result: unknown): string {
  if (!result || !Array.isArray(result)) return "(no references found)";
  if (result.length === 0) return "(no references found)";

  const limit = 30;
  const shown = result.slice(0, limit);
  const formatted = shown
    .map((loc: any, i: number) => {
      const file = uriToPath(loc.uri);
      const range = loc.range;
      return `${i + 1}. ${file} (${range.start.line + 1}:${range.start.character + 1})`;
    })
    .join("\n");

  if (result.length > limit) {
    return `${formatted}\n\n... and ${result.length - limit} more references (showing first ${limit})`;
  }
  return formatted;
}

// Format Completions - Turn completion items into readable text.
function formatCompletions(result: unknown): string {
  if (!result) return "(no completions)";

  // Resolve to CompletionItem[] if it's a CompletionList
  let items: any[];
  if (Array.isArray(result)) {
    items = result;
  } else if (
    result &&
    typeof result === "object" &&
    "items" in result &&
    Array.isArray((result as any).items)
  ) {
    items = (result as any).items;
  } else {
    return JSON.stringify(result, null, 2);
  }

  if (items.length === 0) return "(no completions)";

  // Limit to top 30 to avoid flooding context
  const limited = items.slice(0, 30);
  const lines = limited.map((item: any) => {
    let label = item.label;
    if (item.kind) {
      const kindNames: Record<number, string> = {
        1: "text",
        2: "method",
        3: "function",
        4: "constructor",
        5: "field",
        6: "variable",
        7: "class",
        8: "interface",
        9: "module",
        10: "property",
        11: "unit",
        12: "value",
        13: "enum",
        14: "keyword",
        15: "snippet",
        16: "color",
        17: "file",
        18: "reference",
        19: "folder",
        20: "enumMember",
        21: "constant",
        22: "struct",
        23: "event",
        24: "operator",
        25: "typeParameter",
      };
      label += ` (${kindNames[item.kind] ?? item.kind})`;
    }
    if (item.detail) label += ` — ${item.detail}`;
    return `• ${label}`;
  });

  if (items.length > 30) {
    lines.push(`\n... and ${items.length - 30} more`);
  }
  return lines.join("\n");
}

// Format Document Symbols - Turn symbols into a tree outline.
function formatDocumentSymbols(result: unknown): string {
  if (!result || !Array.isArray(result)) return "(no symbols)";
  const symbols = result as any[];
  if (symbols.length === 0) return "(no symbols)";

  const kindNames: Record<number, string> = {
    1: "File",
    2: "Module",
    3: "Namespace",
    4: "Package",
    5: "Class",
    6: "Method",
    7: "Property",
    8: "Field",
    9: "Constructor",
    10: "Enum",
    11: "Interface",
    12: "Function",
    13: "Variable",
    14: "Constant",
    15: "String",
    16: "Number",
    17: "Boolean",
    18: "Array",
    19: "Object",
    20: "Key",
    21: "Null",
    22: "EnumMember",
    23: "Struct",
    24: "Event",
    25: "Operator",
    26: "TypeParameter",
  };

  function renderSymbol(sym: any, indent = 0): string {
    const pad = "  ".repeat(indent);
    const kind = kindNames[sym.kind] ?? `kind:${sym.kind}`;
    const line = sym.range?.start?.line ? sym.range.start.line + 1 : "?";
    let text = `${pad}${sym.name} [${kind}] :${line}`;
    if (sym.children && Array.isArray(sym.children)) {
      for (const child of sym.children) {
        text += "\n" + renderSymbol(child, indent + 1);
      }
    }
    return text;
  }

  return symbols.map((s) => renderSymbol(s)).join("\n");
}

// Format Diagnostics - Turn diagnostic messages into readable text.
// Format Single Server Diagnostics - Renders one server's diagnostics list.
function formatServerDiagnostics(diags: any[], limit: number): string {
  const severityNames: Record<number, string> = {
    1: "Error",
    2: "Warning",
    3: "Info",
    4: "Hint",
  };

  // Sort By Severity - Errors first, then warnings, info, hints. Ensures
  // the most actionable issues survive truncation.
  diags.sort((a: any, b: any) => (a.severity ?? 99) - (b.severity ?? 99));

  const shown = diags.slice(0, limit);
  const formatted = shown
    .map((d: any, i: number) => {
      const sev = severityNames[d.severity] ?? `sev:${d.severity}`;
      const range = d.range;
      const line = range?.start?.line != null ? range.start.line + 1 : "?";
      const col =
        range?.start?.character != null ? range.start.character + 1 : "?";
      return `${i + 1}. [${sev}] ${d.message} (line ${line}, col ${col})`;
    })
    .join("\n");

  if (diags.length > limit) {
    return `${formatted}\n\n... and ${diags.length - limit} more (showing first ${limit})`;
  }
  return formatted;
}

// Format Diagnostics - Handles the grouped result map from the daemon:
// { [serverId]: { uri, diagnostics[] } }. Single-server results omit the
// header to avoid noise.
function formatDiagnostics(result: unknown, limit = 20): string {
  if (!result || typeof result !== "object") return "(no diagnostics)";

  const grouped = result as Record<string, any>;
  const serverIds = Object.keys(grouped);
  if (serverIds.length === 0) return "(no diagnostics)";

  // Collect Servers With Diagnostics
  const sections: { id: string; diags: any[] }[] = [];
  for (const id of serverIds) {
    const entry = grouped[id];
    const diags = entry?.diagnostics;
    if (Array.isArray(diags) && diags.length > 0) {
      sections.push({ id, diags });
    }
  }

  if (sections.length === 0) return "(no diagnostics)";

  // Single Server - Skip header for brevity.
  if (sections.length === 1) {
    return formatServerDiagnostics(sections[0].diags, limit);
  }

  // Multiple Servers - Group with headers.
  const perServer = Math.max(5, Math.floor(limit / sections.length));
  return sections
    .map((s) => `## ${s.id}\n${formatServerDiagnostics(s.diags, perServer)}`)
    .join("\n\n");
}

// Is Expected Error - Returns true if the error is an expected condition
// (unsupported file type or missing server binary) that should be
// suppressed rather than surfaced to the user.
function isExpectedError(error: unknown): boolean {
  return (
    error instanceof UnsupportedExtensionError ||
    error instanceof ServerNotFoundError
  );
}

// Disabled Servers - In-memory set of disabled server IDs. Scoped to this
// extension process; the shared daemon is never mutated by disable/enable.
const lspToolNames = [
  "lsp_hover",
  "lsp_definition",
  "lsp_references",
  "lsp_completion",
  "lsp_documentSymbol",
  "lsp_diagnostics",
];
const disabledServers = new Set<string>();

// Run LSP Request - Forwards to the daemon, which owns the long-lived
// LspClient cache and handles didOpen/didChange syncing. The daemon
// injects textDocument.uri from the file path, so we omit it here.
// Gated by disabledServers — throws early if the target server is disabled.
async function runLsp(
  filePath: string,
  method: string,
  params: Record<string, unknown>,
): Promise<unknown> {
  try {
    // Check Disabled - The server for this file is blocked; bail before
    // touching the daemon so other pi instances sharing it are unaffected.
    const server = pickServer(filePath);
    if (disabledServers.has(server.id)) {
      throw new Error(
        `LSP server "${server.id}" is disabled. Use /lsp-enable ${server.id} to re-enable.`,
      );
    }
    return await daemonRequest(filePath, server.id, method, params);
  } catch (error) {
    if (isExpectedError(error)) {
      return undefined;
    }
    // Daemon-wrapped errors (plain Error with expected message) are also
    // expected — the daemon catches pickServer() throws and returns them
    // as string error messages.
    if (
      error instanceof Error &&
      (error.message.includes("No LSP server registered") ||
        error.message.includes("not found on PATH"))
    ) {
      return undefined;
    }
    throw error;
  }
}

// Pick Diagnostic Servers - Returns all available, non-disabled servers
// matching the file's extension. Used for fan-out diagnostics.
function pickDiagnosticServers(filePath: string): string[] {
  const ext = path.extname(filePath).replace(/^\./, "");
  return servers
    .filter((s) => s.match.includes(ext) && isServerAvailable(s) && !disabledServers.has(s.id))
    .map((s) => s.id);
}

// Run LSP Diagnostics - Fans out to all matching servers in a single
// daemon call. Returns the grouped result map or undefined if no servers.
async function runDiagnostics(filePath: string): Promise<unknown> {
  try {
    const serverIds = pickDiagnosticServers(filePath);
    if (serverIds.length === 0) return undefined;
    return await daemonDiagnostics(filePath, serverIds, 1500);
  } catch (error) {
    if (isExpectedError(error)) {
      return undefined;
    }
    if (
      error instanceof Error &&
      (error.message.includes("No LSP server registered") ||
        error.message.includes("not found on PATH"))
    ) {
      return undefined;
    }
    throw error;
  }
}

// Shared Parameters Schema - All position-based tools accept file + optional
// position. The "file" field is required; position defaults to (0, 0).
const PositionSchema = Type.Object({
  file: Type.String({ description: "Path to the file to query" }),
  line: Type.Optional(
    Type.Integer({ description: "0-indexed line number (default: 0)" }),
  ),
  character: Type.Optional(
    Type.Integer({ description: "0-indexed character offset (default: 0)" }),
  ),
});

export default function (pi: ExtensionAPI) {
  // Hover - Get hover documentation for the symbol at a position.
  pi.registerTool({
    name: "lsp_hover",
    label: "LSP Hover",
    description: "Get hover documentation for a symbol at a position in a file",
    promptSnippet: "Query LSP server for hover docs at a cursor position",
    parameters: PositionSchema,
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const filePath = path.resolve(ctx.cwd, params.file);
      const lspParams = {
        position: { line: params.line ?? 0, character: params.character ?? 0 },
      };
      const result = await runLsp(filePath, "textDocument/hover", lspParams);
      return {
        content: [{ type: "text", text: formatHover(result) }],
        details: { raw: result },
      };
    },
  });

  // Definition - Find the definition location of a symbol.
  pi.registerTool({
    name: "lsp_definition",
    label: "LSP Definition",
    description: "Find the definition of a symbol at a position in a file",
    promptSnippet: "Query LSP server for symbol definition locations",
    parameters: PositionSchema,
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const filePath = path.resolve(ctx.cwd, params.file);
      const lspParams = {
        position: { line: params.line ?? 0, character: params.character ?? 0 },
      };
      const result = await runLsp(
        filePath,
        "textDocument/definition",
        lspParams,
      );
      return {
        content: [{ type: "text", text: formatDefinition(result) }],
        details: { raw: result },
      };
    },
  });

  // References - Find all references to a symbol.
  pi.registerTool({
    name: "lsp_references",
    label: "LSP References",
    description: "Find all references to a symbol at a position in a file",
    promptSnippet: "Query LSP server for all references to a symbol",
    parameters: PositionSchema,
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const filePath = path.resolve(ctx.cwd, params.file);
      const lspParams = {
        position: { line: params.line ?? 0, character: params.character ?? 0 },
        context: { includeDeclaration: true },
      };
      const result = await runLsp(
        filePath,
        "textDocument/references",
        lspParams,
      );
      return {
        content: [{ type: "text", text: formatReferences(result) }],
        details: { raw: result },
      };
    },
  });

  // Completion - Get completion suggestions at a position.
  pi.registerTool({
    name: "lsp_completion",
    label: "LSP Completion",
    description: "Get completion suggestions at a position in a file",
    promptSnippet: "Query LSP server for completion items at a cursor position",
    parameters: PositionSchema,
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const filePath = path.resolve(ctx.cwd, params.file);
      const lspParams = {
        position: { line: params.line ?? 0, character: params.character ?? 0 },
      };
      const result = await runLsp(
        filePath,
        "textDocument/completion",
        lspParams,
      );
      return {
        content: [{ type: "text", text: formatCompletions(result) }],
        details: { raw: result },
      };
    },
  });

  // Document Symbol - Get the outline/symbol tree of a file.
  pi.registerTool({
    name: "lsp_documentSymbol",
    label: "LSP Document Symbols",
    description: "Get the symbol outline (classes, functions, etc.) of a file",
    promptSnippet: "Query LSP server for document symbol outline",
    parameters: Type.Object({
      file: Type.String({ description: "Path to the file" }),
    }),
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const filePath = path.resolve(ctx.cwd, params.file);
      const result = await runLsp(
        filePath,
        "textDocument/documentSymbol",
        {},
      );
      return {
        content: [{ type: "text", text: formatDocumentSymbols(result) }],
        details: { raw: result },
      };
    },
  });

  // Diagnostics - Get lint/type-check diagnostics for a file.
  pi.registerTool({
    name: "lsp_diagnostics",
    label: "LSP Diagnostics",
    description:
      "Get diagnostics (errors, warnings) for a file from its LSP server",
    promptSnippet: "Query LSP server for file diagnostics (errors, warnings)",
    parameters: Type.Object({
      file: Type.String({ description: "Path to the file" }),
    }),
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const filePath = path.resolve(ctx.cwd, params.file);
      const result = await runDiagnostics(filePath);
      return {
        content: [{ type: "text", text: formatDiagnostics(result) }],
        details: { raw: result },
      };
    },
  });

  // Auto-Check Flag - Enable automatic diagnostics after edit/write
  pi.registerFlag("lsp-auto-check", {
    description:
      "Automatically run LSP diagnostics after edit/write operations",
    type: "boolean",
    default: true,
  });

  // Background Init on Read - Fire-and-forget LSP initialization when the
  // LLM reads a file with a supported extension. Doesn't block the read,
  // just ensures the server is warm by the time an LSP tool is called.
  pi.on("tool_result", async (event, ctx) => {
    if (event.toolName !== "read" || event.isError) return;

    const filePath = event.input?.path;
    if (!filePath || typeof filePath !== "string") return;

    try {
      const absolutePath = path.resolve(ctx.cwd, filePath);
      // Warm Diagnostic Servers - Fire-and-forget so servers are ready by
      // the time an LSP tool is called.
      const serverIds = pickDiagnosticServers(absolutePath);
      if (serverIds.length > 0) {
        void daemonDiagnostics(absolutePath, serverIds).catch(() => {});
      }
    } catch {
      // Silently ignore — unsupported file type, missing binary, etc.
    }
  });

  // Auto-Check After Edit/Write - Run diagnostics automatically
  pi.on("tool_result", async (event, ctx) => {
    // Check Enabled
    if (!pi.getFlag("lsp-auto-check")) return;

    // Skip If All Disabled - No LSP server is available for this instance.
    if (servers.every((s) => disabledServers.has(s.id))) return;

    // Edit & Write Only
    if (!["edit", "write"].includes(event.toolName)) return;

    // Skip Errors
    if (event.isError) return;

    // Resolve Filepath
    const filePath = event.input?.path;
    if (!filePath || typeof filePath !== "string") return;
    const absolutePath = path.resolve(ctx.cwd, filePath);

    try {
      // Run LSP diagnostics
      const result = await runDiagnostics(absolutePath);
      const formatted = formatDiagnostics(result, 10);

      // Only send a message if there are actual diagnostics
      if (formatted !== "(no diagnostics)") {
        pi.sendMessage(
          {
            customType: "lsp-auto-check",
            content: `LSP found issues in ${filePath}:\n${formatted}`,
            display: true,
            details: { filePath, diagnostics: result },
          },
          { deliverAs: "steer" },
        );
      }
    } catch (error) {
      // Silently fail - don't interrupt the flow
      // Only log if there's an actual error we care about
      if (!isExpectedError(error)) {
        const msg =
          error && typeof error === "object" && "message" in error
            ? (error as { message: string }).message
            : String(error);
        console.error("LSP auto-check failed:", msg);
      }
    }
  });

  // Manual Check Command - Run diagnostics on specific files
  pi.registerCommand("lsp-check", {
    description: "Run LSP diagnostics on one or more files",
    handler: async (args, ctx) => {
      const files = args ? args.split(/\s+/) : [];

      if (files.length === 0) {
        ctx.ui.notify("Usage: /lsp-check <file1> [file2...]", "error");
        return;
      }

      for (const file of files) {
        const absolutePath = path.resolve(ctx.cwd, file);
        try {
          const result = await runDiagnostics(absolutePath);
          const formatted = formatDiagnostics(result);

          if (formatted === "(no diagnostics)") {
            ctx.ui.notify(`${file}: ✓ No issues`, "info");
          } else {
            ctx.ui.notify(`${file}:\n${formatted}`, "warning");
          }
        } catch (error) {
          const msg =
            error && typeof error === "object" && "message" in error
              ? (error as { message: string }).message
              : "Unknown error";
          ctx.ui.notify(`${file}: Failed to check - ${msg}`, "error");
        }
      }
    },
  });

  // --- Server Control Commands ---

  // Shared Argument Completions - Suggests registered server IDs plus "all".
  const serverCompletions = (prefix: string) => {
    const ids = [...servers.map((s) => s.id), "all"].filter((id) =>
      id.startsWith(prefix),
    );
    return ids.length > 0 ? ids.map((id) => ({ value: id, label: id })) : null;
  };

  // Parse Server IDs - Validates args against registered servers. Bare/empty
  // or "all" returns every server ID. Throws on unknown names.
  function parseServerIds(args: string | undefined): string[] {
    if (!args || !args.trim()) return servers.map((s) => s.id);
    const ids = args.trim().split(/\s+/);
    if (ids.includes("all")) return servers.map((s) => s.id);
    const invalid = ids.filter((id) => !servers.some((s) => s.id === id));
    if (invalid.length > 0) {
      throw new Error(
        `Unknown server(s): ${invalid.join(
          ", ",
        )}. Available: ${servers.map((s) => s.id).join(", ")}`,
      );
    }
    return ids;
  }

  // Update Tool Visibility - When all servers are disabled, remove LSP tools
  // from the active set so the LLM won't attempt them. When any is enabled,
  // restore them. Captures current active tools at toggle time.
  function updateToolVisibility(): void {
    const current = pi.getActiveTools();
    if (servers.every((s) => disabledServers.has(s.id))) {
      // All disabled — strip LSP tools
      pi.setActiveTools(current.filter((name) => !lspToolNames.includes(name)));
    } else {
      // Any enabled — merge LSP tools back in
      const merged = [...new Set([...current, ...lspToolNames])];
      pi.setActiveTools(merged);
    }
  }

  // List Servers - Show running daemon entries and disabled state.
  pi.registerCommand("lsp-servers", {
    description: "List running LSP servers and disabled state",
    handler: async (_args, ctx) => {
      try {
        const status = (await daemonStatus()) as { servers?: unknown[] };
        const running = Array.isArray(status?.servers) ? status.servers : [];
        const disabled = Array.from(disabledServers);

        if (running.length === 0 && disabled.length === 0) {
          ctx.ui.notify("No running servers, none disabled", "info");
          return;
        }

        let msg = "";
        if (running.length > 0) {
          msg += `Running: ${running.map((s: any) => s.id).join(", ")}\n`;
        }
        if (disabled.length > 0) {
          msg += `Disabled: ${disabled.join(", ")}`;
        }
        ctx.ui.notify(msg.trim(), "info");
      } catch (error) {
        const msg =
          error && typeof error === "object" && "message" in error
            ? (error as { message: string }).message
            : "Unknown error";
        ctx.ui.notify(`Status failed: ${msg}`, "error");
      }
    },
  });

  // Disable Servers - Add to disabled set; removes LSP tools when all are
  // disabled so the LLM won't waste context on them.
  pi.registerCommand("lsp-disable", {
    description:
      "Disable LSP server(s) — bare command disables all. Removes tools when all are disabled.",
    getArgumentCompletions: serverCompletions,
    handler: async (args, ctx) => {
      try {
        const ids = parseServerIds(args);
        for (const id of ids) {
          disabledServers.add(id);
        }
        updateToolVisibility();
        const label =
          ids.length === servers.length ? "all servers" : ids.join(", ");
        ctx.ui.notify(`Disabled: ${label}`, "info");
      } catch (error) {
        const msg =
          error && typeof error === "object" && "message" in error
            ? (error as { message: string }).message
            : "Unknown error";
        ctx.ui.notify(msg, "error");
      }
    },
  });

  // Enable Servers - Remove from disabled set; restores LSP tools when any
  // server becomes available.
  pi.registerCommand("lsp-enable", {
    description:
      "Enable LSP server(s) — bare command enables all. Restores tools when any is enabled.",
    getArgumentCompletions: serverCompletions,
    handler: async (args, ctx) => {
      try {
        const ids = parseServerIds(args);
        for (const id of ids) {
          disabledServers.delete(id);
        }
        updateToolVisibility();
        const label =
          ids.length === servers.length ? "all servers" : ids.join(", ");
        ctx.ui.notify(`Enabled: ${label}`, "info");
      } catch (error) {
        const msg =
          error && typeof error === "object" && "message" in error
            ? (error as { message: string }).message
            : "Unknown error";
        ctx.ui.notify(msg, "error");
      }
    },
  });

  // Destroy Servers - Kill running LspClient entries in the daemon. Entries
  // can respawn on next request; pair with /lsp-disable to also block.
  pi.registerCommand("lsp-destroy", {
    description:
      "Kill running LSP server process(es) in the daemon — bare command destroys all.",
    getArgumentCompletions: serverCompletions,
    handler: async (args, ctx) => {
      try {
        const ids = parseServerIds(args);
        if (ids.length === servers.length) {
          await daemonDestroyServer();
        } else {
          for (const id of ids) {
            await daemonDestroyServer(id);
          }
        }
        const label =
          ids.length === servers.length ? "all servers" : ids.join(", ");
        ctx.ui.notify(`Destroyed: ${label}`, "info");
      } catch (error) {
        const msg =
          error && typeof error === "object" && "message" in error
            ? (error as { message: string }).message
            : "Unknown error";
        ctx.ui.notify(msg, "error");
      }
    },
  });
}