/**
 * Codexis - Code index query tool for pi
 *
 * Provides a single tool that queries the .codexis/index.db SQLite database
 * containing symbols, files, and line numbers for the codebase.
 *
 * - Only registers if `codexis` binary is in PATH
 * - Auto-indexes on first tool call if DB is missing
 * - Re-indexes (incremental) on each call to keep index fresh
 */

import { readFileSync } from "node:fs";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import { Type } from "@sinclair/typebox";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { existsSync } from "node:fs";
import Database from "better-sqlite3";

const baseDir = dirname(fileURLToPath(import.meta.url));

const SCHEMA = readFileSync(join(baseDir, "schema.sql"), "utf-8");

const DESCRIPTION = `Query the code index database (.codexis/index.db). Run read-only SQL to find symbols, files, and line numbers across the codebase.

${SCHEMA}

Example queries:
  -- Find where a function is defined
  SELECT f.path, s.line FROM symbols s JOIN files f ON s.file_id=f.id WHERE s.name='HandleRequest'

  -- Public API of a package
  SELECT s.name, s.kind, s.line, f.path FROM symbols s JOIN files f ON s.file_id=f.id WHERE f.package='server' AND s.exported=1

  -- All types in a directory
  SELECT s.name, s.line, f.path FROM symbols s JOIN files f ON s.file_id=f.id WHERE f.path LIKE 'backend/api/%' AND s.kind='type'

  -- Methods on a class/type (via parent_id)
  SELECT c.name as parent, s.name, s.kind, s.line FROM symbols s JOIN symbols c ON s.parent_id=c.id WHERE c.name='AuthService'

  -- Overview: symbols per area
  SELECT CASE WHEN f.path LIKE 'backend/%' THEN 'backend' WHEN f.path LIKE 'frontend/%' THEN 'frontend' ELSE 'other' END as area, COUNT(*) FROM symbols s JOIN files f ON s.file_id=f.id GROUP BY area

  -- Full-text search for content across all files
  SELECT f.path, snippet(file_contents, 1, '>>>', '<<<', '...', 20) as match FROM file_contents fc JOIN files f ON f.id=fc.file_id WHERE file_contents MATCH 'handleRequest' ORDER BY rank LIMIT 10

  -- FTS search scoped to a directory
  SELECT f.path, snippet(file_contents, 1, '>>>', '<<<', '...', 20) as match FROM file_contents fc JOIN files f ON f.id=fc.file_id WHERE file_contents MATCH 'database migration' AND f.path LIKE 'backend/%' ORDER BY rank LIMIT 10`;

function findGitRoot(cwd: string): string | null {
  try {
    const { execSync } = require("node:child_process");
    return execSync("git rev-parse --show-toplevel", {
      cwd,
      encoding: "utf-8",
      stdio: ["pipe", "pipe", "pipe"],
    }).trim();
  } catch {
    return null;
  }
}

function codexisAvailable(): boolean {
  try {
    const { execSync } = require("node:child_process");
    execSync("which codexis", { stdio: ["pipe", "pipe", "pipe"] });
    return true;
  } catch {
    return false;
  }
}

export default function (pi: ExtensionAPI) {
  if (!codexisAvailable()) {
    return;
  }

  pi.registerTool({
    name: "Codexis",
    label: "Codexis",
    description: DESCRIPTION,
    parameters: Type.Object({
      type: Type.Enum(
        {
          sql: "sql",
          call_graph: "call_graph",
        },
        {
          description:
            "Query type: 'sql' for custom SQL queries, 'call_graph' for function call graph",
        },
      ),
      sql: Type.Optional(
        Type.String({
          description:
            "SQL query to run (required when type='sql', ignored when type='call_graph')",
        }),
      ),
    }),

    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
      const gitRoot = findGitRoot(ctx.cwd);
      if (!gitRoot) {
        throw new Error("Not in a git repository");
      }

      const dbPath = join(gitRoot, ".codexis", "index.db");

      // Run incremental index to keep DB fresh (fast if nothing changed)
      const indexResult = await pi.exec("codexis", [gitRoot], {
        signal,
        timeout: 120_000,
      });

      if (indexResult.code !== 0) {
        throw new Error(
          `codexis indexing failed (exit ${indexResult.code}): ${indexResult.stderr}`,
        );
      }

      if (!existsSync(dbPath)) {
        throw new Error(
          "codexis ran but no index.db was produced. Check codexis output.",
        );
      }

      const db = new Database(dbPath, { readonly: true });
      try {
        let rows: Record<string, unknown>[];

        if (params.type === "call_graph") {
          const callGraphSql = `
            WITH defs AS (
              SELECT s.id, s.name, s.line as start_line, COALESCE(s.line_end, s.line + 1) as end_line,
                     f.path, f.id as file_id
              FROM symbols s JOIN files f ON s.file_id = f.id
              WHERE s.kind IN ('function', 'method') AND s.parent_id IS NULL
            ),
            func_counts AS (
              SELECT name, COUNT(*) as n FROM defs GROUP BY name
            ),
            calls AS (
              SELECT s.name, s.line as call_line, f.id as caller_file_id
              FROM symbols s JOIN files f ON s.file_id = f.id
              WHERE s.kind = 'reference'
            ),
            call_sites AS (
              SELECT DISTINCT
                caller.name as caller,
                caller.path as caller_file,
                caller.start_line || '-' || caller.end_line as caller_lines,
                c.call_line,
                c.name as callee_name,
                fc.n as defs_count
              FROM calls c
              JOIN defs caller ON c.caller_file_id = caller.file_id 
                AND c.call_line >= caller.start_line AND c.call_line <= caller.end_line
              JOIN func_counts fc ON c.name = fc.name
            )
            SELECT 
              caller,
              caller_file,
              caller_lines,
              call_line,
              callee_name as callee,
              CASE WHEN defs_count > 1 THEN '' ELSE callee.path END as callee_file,
              CASE WHEN defs_count > 1 THEN '' ELSE callee.start_line || '-' || callee.end_line END as callee_lines
            FROM call_sites
            LEFT JOIN defs callee ON call_sites.callee_name = callee.name AND defs_count <= 1
            ORDER BY caller_file, call_line;
          `;
          rows = db.prepare(callGraphSql).all();
        } else {
          // type === 'sql' mode
          if (!params.sql) {
            throw new Error("sql parameter is required when type='sql'");
          }
          const normalized = params.sql.trim().toUpperCase();
          if (
            !normalized.startsWith("SELECT") &&
            !normalized.startsWith("WITH") &&
            !normalized.startsWith("EXPLAIN") &&
            !normalized.startsWith("PRAGMA")
          ) {
            throw new Error(
              "Only SELECT, WITH, EXPLAIN, and PRAGMA queries are allowed",
            );
          }
          rows = db.prepare(params.sql).all();
        }

        if (rows.length === 0) {
          return {
            content: [{ type: "text", text: "No results." }],
            details: { rowCount: 0 },
          };
        }

        if (params.type === "call_graph") {
          // Format call graph with file aliases
          const fileMap = new Map<string, string>();
          let aliasCounter = 1;
          const fileToAlias = (path: string) => {
            if (!fileMap.has(path)) {
              fileMap.set(path, String.fromCharCode(64 + aliasCounter++)); // A, B, C...
            }
            return fileMap.get(path)!;
          };

          // First pass: collect all files
          for (const row of rows) {
            const r = row as Record<string, unknown>;
            fileToAlias(String(r.caller_file));
            if (r.callee_file) {
              fileToAlias(String(r.callee_file));
            }
          }

          // Build file legend
          const fileLegend = Array.from(fileMap.entries())
            .map(([path, alias]) => `  ${alias} = ${path}`)
            .join("\n");

          // Group calls by caller
          const grouped = new Map<string, Array<Record<string, unknown>>>();
          for (const row of rows) {
            const key = `${row.caller} (${row.caller_file}:${row.caller_lines})`;
            if (!grouped.has(key)) {
              grouped.set(key, []);
            }
            grouped.get(key)!.push(row);
          }

          // Build call graph output
          const callGraphLines: string[] = [];
          callGraphLines.push("Files:");
          callGraphLines.push(fileLegend);
          callGraphLines.push("");

          for (const [callerKey, calls] of grouped) {
            callGraphLines.push(callerKey);
            for (const call of calls) {
              const r = call as Record<string, unknown>;
              const callLine = String(r.call_line);
              const callee = String(r.callee);
              const calleeFile = String(r.callee_file);
              const calleeLines = String(r.callee_lines);
              const calleeAlias = calleeFile ? fileToAlias(calleeFile) : "??";
              const target = calleeFile
                ? `${calleeAlias}:${calleeLines}`
                : "?? (ambiguous)";
              callGraphLines.push(
                `  ${callLine.padEnd(4)} -> ${callee.padEnd(20)} ${target}`,
              );
            }
          }

          const result = callGraphLines.join("\n");
          return {
            content: [{ type: "text", text: result }],
            details: { rowCount: rows.length },
          };
        }

        // Default: Format as aligned text table
        const columns = Object.keys(rows[0] as Record<string, unknown>);
        const data = rows.map((row) => {
          const r = row as Record<string, unknown>;
          return columns.map((col) => String(r[col] ?? "NULL"));
        });

        const widths = columns.map((col, i) =>
          Math.max(col.length, ...data.map((row) => row[i].length)),
        );

        const header = columns
          .map((col, i) => col.padEnd(widths[i]))
          .join("  ");
        const separator = widths.map((w) => "-".repeat(w)).join("  ");
        const body = data
          .map((row) => row.map((val, i) => val.padEnd(widths[i])).join("  "))
          .join("\n");

        const result = `${header}\n${separator}\n${body}`;

        // Truncate if huge
        const maxLen = 48000;
        const truncated =
          result.length > maxLen
            ? result.slice(0, maxLen) +
              `\n\n[Truncated: ${rows.length} rows total, showing partial results. Narrow your query.]`
            : result;

        return {
          content: [{ type: "text", text: truncated }],
          details: { rowCount: rows.length },
        };
      } finally {
        db.close();
      }
    },
  });
}