/**
 * pi-julia — Pi extension
 *
 * Two complementary Julia services, both started on demand:
 *
 *   DaemonMode server  — warm Julia process for TTFX-free code execution
 *   LanguageServer.jl  — full LSP for hover docs, diagnostics, completions, go-to-def
 *
 * User commands:
 *   /julia-start [path]    — start both servers for a Julia project
 *   /julia-stop            — stop both servers
 *   /julia-activate [path] — Julia-flavoured alias for /julia-start
 *
 * LLM tools:
 *   julia_start(project_path)         — start (or restart) both servers
 *   julia_stop()                      — stop both servers
 *   julia_run_file(file, args?)       — execute a .jl file (DaemonMode)
 *   julia_run_expr(code)              — execute a Julia snippet (DaemonMode)
 *   julia_doc(symbol)                 — docstring for a symbol by name (LSP)
 *   julia_hover(file, line, col)      — docstring + type at position (LSP)
 *   julia_diagnostics(file?)          — errors and warnings (LSP)
 *   julia_complete(file, line, col)   — completions at position (LSP)
 *   julia_definition(file, line, col) — go-to-definition (LSP)
 *
 * Settings (in ~/.pi/agent/settings.json under "pi-julia"):
 *   projectPath  — remembered project path
 *   autoStart    — auto-start on session_start (default: false)
 */

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

import { DaemonManager, ensureDaemonInstalled } from "./daemon.ts";
import {
  LspClient,
  spawnLsp,
  findLsEnv,
  formatDiagnostics,
  formatAllDiagnostics,
} from "./lsp.ts";
import {
  findProjectToml,
  resolveProjectPath,
  formatOutput,
} from "./utils.ts";

// ─── Settings ─────────────────────────────────────────────────────────────────

const SETTINGS_PATH = path.join(os.homedir(), ".pi", "agent", "settings.json");

interface PiJuliaSettings {
  projectPath?: string;
  autoStart?: boolean;
}

function readPiJuliaSettings(): PiJuliaSettings {
  try {
    const all = JSON.parse(fs.readFileSync(SETTINGS_PATH, "utf-8"));
    const s = all["pi-julia"];
    return s && typeof s === "object" ? (s as PiJuliaSettings) : {};
  } catch { return {}; }
}

function writePiJuliaSettings(patch: Partial<PiJuliaSettings>): void {
  let all: Record<string, unknown> = {};
  try { all = JSON.parse(fs.readFileSync(SETTINGS_PATH, "utf-8")); } catch {}
  all["pi-julia"] = { ...(all["pi-julia"] as object ?? {}), ...patch };
  fs.writeFileSync(SETTINGS_PATH, JSON.stringify(all, null, 2) + "\n");
}

// ─── Server state ─────────────────────────────────────────────────────────────

const dm = new DaemonManager();

let lsp: LspClient | undefined;
let lspReady      = false;
let agentTurnCount = 0;
let lspStarting = false;

function stopLsp(): void {
  if (lsp) {
    try { lsp.shutdown(); } catch {}
    lsp         = undefined;
    lspReady    = false;
    lspStarting = false;
  }
}

// ─── LSP background start ─────────────────────────────────────────────────────

// Launch LSP in the background — does not block startBoth().
// Updates status widget when ready.
async function startLspBackground(
  pi: ExtensionAPI,
  ctx: ExtensionContext,
  projectPath: string | undefined
): Promise<void> {
  if (lspStarting) return;
  lspStarting = true;
  lspReady    = false;

  const name = projectPath ? path.basename(projectPath) : "global";
  const cwd  = projectPath ?? ctx.cwd;

  const execFn = (cmd: string, args: string[], opts: { timeout: number }) =>
    pi.exec(cmd, args, { ...opts, cwd });

  const lsEnvPath = await findLsEnv(execFn);
  if (lsEnvPath === null) {
    lspStarting = false;
    updateStatus(ctx, name);
    ctx.ui.notify("pi-julia: could not find or install LanguageServer.jl", "warning");
    return;
  }

  stopLsp();
  lsp = spawnLsp(projectPath ?? "", lsEnvPath);

  try {
    await lsp.initialize(projectPath); // up to 90s
    lspReady    = true;
    lspStarting = false;
    updateStatus(ctx, name);
  } catch (e) {
    lspStarting = false;
    lspReady    = false;
    updateStatus(ctx, name);
    ctx.ui.notify(`pi-julia: LSP failed to initialize — ${e}`, "warning");
  }
}

// ─── Status ───────────────────────────────────────────────────────────────────

function updateStatus(ctx: ExtensionContext, projectName: string): void {
  const dmStr = dm.isRunning ? `:${dm.port}` : "off";
  const lsStr = lspStarting ? "LSP:starting…"
              : lspReady    ? "LSP:ready"
              :               "LSP:off";
  ctx.ui.setStatus("pi-julia", `Julia[${projectName}] DM:${dmStr} ${lsStr}`);
}

// ─── LSP doc helper ───────────────────────────────────────────────────────────

// Open a file in the LSP if not already open, then wait until the server
// publishes its first diagnostics for it — that notification is the signal
// that the LS has finished its initial parse. Timeouts after 30s.
async function ensureDocOpen(filePath: string, timeoutMs = 30_000): Promise<void> {
  if (lsp!.isOpen(filePath)) return;
  lsp!.openDoc(filePath);
  const deadline = Date.now() + timeoutMs;
  while (!lsp!.hasDiagnosticsFor(filePath) && Date.now() < deadline) {
    await new Promise((r) => setTimeout(r, 300));
  }
}

// ─── Extension ────────────────────────────────────────────────────────────────

export default async function (pi: ExtensionAPI): Promise<void> {

  // ── Core: start both servers ────────────────────────────────────────────

  async function startBoth(
    rawPath: string,
    pi: ExtensionAPI,
    ctx: ExtensionContext
  ): Promise<string> {
    const resolved = resolveProjectPath(rawPath, ctx.cwd);
    if (!resolved && rawPath.trim())
      return `No Project.toml at: ${path.resolve(ctx.cwd, rawPath)}`;
    const projectPath = resolved; // undefined = global env

    writePiJuliaSettings({ projectPath });
    const name = projectPath ? path.basename(projectPath) : "global";
    ctx.ui.setStatus("pi-julia", `Julia[${name}]: starting…`);

    // ── DaemonMode (blocking — wait until port is open) ─────────────────
    const execFn = (cmd: string, args: string[], opts: { timeout: number }) =>
      pi.exec(cmd, args, { ...opts, cwd: projectPath ?? ctx.cwd });
    const dmOk = await ensureDaemonInstalled(execFn);
    if (!dmOk) {
      ctx.ui.setStatus("pi-julia", "Julia: DaemonMode install failed");
      return "Could not install DaemonMode.jl — check Julia setup.";
    }
    dm.launch(projectPath);
    const dmReady = await dm.waitUntilReady();
    if (!dmReady) {
      dm.stop();
      return "DaemonMode server timed out. Try /julia-start again.";
    }

    updateStatus(ctx, name);

    // ── LSP (non-blocking — starts in background) ────────────────────────
    startLspBackground(pi, ctx, projectPath); // intentionally not awaited

    return `DaemonMode ready on port ${dm.port}. LSP starting in background…`;
  }

  // ── session_start ───────────────────────────────────────────────────────

  pi.on("session_start", async (_e, ctx) => {
    agentTurnCount = 0;
    const settings = readPiJuliaSettings();
    if (!settings.autoStart) {
      ctx.ui.setStatus("pi-julia", "Julia: stopped");
      return;
    }
    const savedPath = settings.projectPath
      ? resolveProjectPath(settings.projectPath, ctx.cwd)
      : findProjectToml(ctx.cwd);
    // savedPath undefined → no project found → fall back to global env, like Julia itself
    await startBoth(savedPath ?? "", pi, ctx);
  });

  // ── session_shutdown ────────────────────────────────────────────────────

  pi.on("session_shutdown", async () => {
    dm.stop();
    stopLsp();
  });

  // ── before_agent_start ──────────────────────────────────────────────────

  pi.on("before_agent_start", async (event, ctx) => {
    const hasJulia =
      fs.existsSync(path.join(ctx.cwd, "Project.toml")) ||
      fs.readdirSync(ctx.cwd).some((f) => f.endsWith(".jl"));
    if (!hasJulia) return;

    const isFirstTurn = agentTurnCount === 0;
    agentTurnCount++;

    const dmStatus = dm.isRunning
      ? `DaemonMode RUNNING on port ${dm.port}.`
      : `DaemonMode NOT running — call \`julia_start\` first.`;
    const lsStatus = lspReady ? "LanguageServer READY."
      : lspStarting            ? "LanguageServer starting…"
      :                          "LanguageServer not running.";

    let hint: string;
    if (isFirstTurn) {
      // Full rules on the first turn — establish the constraint clearly.
      hint =
        `\n\n## Julia — IMPORTANT RULES (enforced by pi-julia)\n` +
        `**NEVER run \`julia\` via bash.** Calling \`julia\` directly wastes ~5 seconds on every invocation.\n` +
        `**ALWAYS use the \`julia_*\` tools below instead.**\n\n` +
        `Current status: ${dmStatus} ${lsStatus}\n\n` +
        `**Execution tools** (require DaemonMode — no TTFX overhead):\n` +
        `- \`julia_start\` — start the server (do this first, once per session)\n` +
        `- \`julia_run_file\` — run a .jl file\n` +
        `- \`julia_run_expr\` — run a Julia expression or code block\n\n` +
        `**Code intelligence tools** (require LanguageServer):\n` +
        `- \`julia_doc\` — docstring for a symbol by name\n` +
        `- \`julia_hover\`, \`julia_diagnostics\`, \`julia_complete\`, \`julia_definition\``;
    } else {
      // Subsequent turns: status only — rules already established.
      hint = `\n\n## Julia (pi-julia) — ${dmStatus} ${lsStatus}`;
    }

    return { systemPrompt: event.systemPrompt + hint };
  });

  // ── /julia-start ────────────────────────────────────────────────────────

  pi.registerCommand("julia-start", {
    description: "Start Julia servers (DaemonMode + LSP). Optionally pass a project path.",
    handler: async (args, ctx) => {
      const msg = await startBoth(args ?? "", pi, ctx);
      ctx.ui.notify(msg, dm.isRunning ? "info" : "error");
    },
  });

  pi.registerCommand("julia-activate", {
    description: "Activate a Julia project. Alias for /julia-start.",
    handler: async (args, ctx) => {
      const msg = await startBoth(args ?? "", pi, ctx);
      ctx.ui.notify(msg, dm.isRunning ? "info" : "error");
    },
  });

  pi.registerCommand("julia-stop", {
    description: "Stop both Julia servers.",
    handler: async (_args, ctx) => {
      if (!dm.isRunning && !lsp) {
        ctx.ui.notify("No Julia servers are running.", "info");
        return;
      }
      const name = dm.projectPath ? path.basename(dm.projectPath) : "global";
      dm.stop();
      stopLsp();
      ctx.ui.setStatus("pi-julia", "Julia: stopped");
      ctx.ui.notify(`Julia servers for ${name} stopped.`, "info");
    },
  });

  // ── Guards ───────────────────────────────────────────────────────────────

  function daemonGuard(): string | null {
    if (!dm.isRunning)
      return 'No Julia server running. Call julia_start("") first.';
    return null;
  }

  function lspGuard(): string | null {
    if (!lsp)
      return 'Language server not running. Call julia_start("") first.';
    if (lspStarting)
      return "Language server is still starting. Please wait a moment and try again.";
    if (!lspReady)
      return "Language server is not ready. Try julia_start again.";
    return null;
  }

  // ── Tool: julia_start ────────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_start",
    label: "Start Julia servers",
    description:
      "Start DaemonMode (for execution) and LanguageServer.jl (for code intelligence). " +
      'Pass "" or "." to auto-detect the nearest Project.toml from the working directory.',
    parameters: Type.Object({
      project_path: Type.String({
        description: 'Julia project directory. Pass "" or "." to auto-detect.',
      }),
    }),
    async execute(_id, params, _signal, _onUpdate, ctx) {
      const msg = await startBoth(params.project_path, pi, ctx);
      return { content: [{ type: "text", text: msg }], details: { running: dm.isRunning } };
    },
  });

  // ── Tool: julia_stop ─────────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_stop",
    label: "Stop Julia servers",
    description: "Stop both DaemonMode and LanguageServer to free resources.",
    parameters: Type.Object({}),
    async execute(_id, _p, _signal, _onUpdate, ctx) {
      if (!dm.isRunning && !lsp)
        return { content: [{ type: "text", text: "No Julia servers running." }], details: {} };
      const name = dm.projectPath ? path.basename(dm.projectPath) : "global";
      dm.stop(); stopLsp();
      ctx.ui.setStatus("pi-julia", "Julia: stopped");
      return { content: [{ type: "text", text: `Servers for ${name} stopped.` }], details: {} };
    },
  });

  // ── Tool: julia_run_file ─────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_run_file",
    label: "Run Julia file",
    description:
      "Execute a .jl file against the warm DaemonMode server (no TTFX). " +
      "Call julia_start first if not running.",
    parameters: Type.Object({
      file: Type.String({ description: "Path to the .jl file (absolute or relative to cwd)." }),
      args: Type.Optional(Type.Array(Type.String(), { description: "Arguments passed as ARGS." })),
    }),
    async execute(_id, params, signal, _onUpdate, ctx) {
      const g = daemonGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      const filePath = path.resolve(ctx.cwd, params.file);
      if (!fs.existsSync(filePath))
        return { content: [{ type: "text", text: `File not found: ${filePath}` }], details: {} };

      const result = await pi.exec(
        "julia",
        ["--startup-file=no", "-e", `using DaemonMode; runargs(${dm.port})`, filePath, ...(params.args ?? [])],
        { cwd: ctx.cwd, signal, timeout: 300_000 }
      );
      return { content: [{ type: "text", text: formatOutput(result) }], details: { exitCode: result.code } };
    },
  });

  // ── Tool: julia_run_expr ─────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_run_expr",
    label: "Run Julia expression",
    description:
      "Execute a Julia code block against the warm DaemonMode server (no TTFX). " +
      "Call julia_start first if not running.",
    parameters: Type.Object({
      code: Type.String({ description: "Julia code to execute. Can be multi-line." }),
    }),
    async execute(_id, params, signal, _onUpdate, ctx) {
      const g = daemonGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      const tmpFile = path.join(os.tmpdir(), `pi-julia-${process.pid}-${Date.now()}.jl`);
      fs.writeFileSync(tmpFile, params.code, "utf-8");

      try {
        const result = await pi.exec(
          "julia",
          ["--startup-file=no", "-e", `using DaemonMode; runargs(${dm.port})`, tmpFile],
          { cwd: ctx.cwd, signal, timeout: 300_000 }
        );
        return { content: [{ type: "text", text: formatOutput(result) }], details: { exitCode: result.code } };
      } finally {
        try { fs.unlinkSync(tmpFile); } catch {}
      }
    },
  });

  // ── Tool: julia_doc ──────────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_doc",
    label: "Julia documentation",
    description:
      "Fetch the docstring for a Julia symbol by name (function, type, macro, or module). " +
      "Uses the LSP server — no Julia process startup overhead. " +
      'Supports qualified names like "Base.sort" and macros like "@time". ' +
      "Requires the LSP server (started by julia_start).",
    parameters: Type.Object({
      symbol: Type.String({
        description: 'Julia symbol to look up, e.g. "sort", "Base.sort", "push!", "@time".',
      }),
    }),
    async execute(_id, params, _signal, _onUpdate, _ctx) {
      const g = lspGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      const result = await lsp!.docByName(params.symbol).catch((e) => `LSP error: ${e}`);
      return {
        content: [{ type: "text", text: result ?? `No documentation found for \`${params.symbol}\`.` }],
        details: {},
      };
    },
  });

  // ── Tool: julia_hover ────────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_hover",
    label: "Julia hover info",
    description:
      "Get the docstring and type information for the symbol at a given position in a Julia file. " +
      "Requires the LSP server (started by julia_start).",
    parameters: Type.Object({
      file: Type.String({ description: "Path to the .jl file." }),
      line: Type.Number({ description: "Line number (1-based)." }),
      col:  Type.Number({ description: "Column number (1-based)." }),
    }),
    async execute(_id, params, _signal, _onUpdate, ctx) {
      const g = lspGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      const filePath = path.resolve(ctx.cwd, params.file);
      if (!fs.existsSync(filePath))
        return { content: [{ type: "text", text: `File not found: ${filePath}` }], details: {} };

      await ensureDocOpen(filePath);

      const result = await lsp!.hover(filePath, params.line - 1, params.col - 1)
        .catch((e) => `LSP error: ${e}`);

      return {
        content: [{ type: "text", text: result ?? "No hover information available." }],
        details: {},
      };
    },
  });

  // ── Tool: julia_diagnostics ──────────────────────────────────────────────

  pi.registerTool({
    name: "julia_diagnostics",
    label: "Julia diagnostics",
    description:
      "Get errors and warnings for a Julia file (or all open files) from the language server. " +
      "Diagnostics are computed statically — no code is executed. " +
      "Requires the LSP server.",
    parameters: Type.Object({
      file: Type.Optional(Type.String({
        description: "Path to the .jl file. Omit to get diagnostics for all open files.",
      })),
    }),
    async execute(_id, params, _signal, _onUpdate, ctx) {
      const g = lspGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      if (params.file) {
        const filePath = path.resolve(ctx.cwd, params.file);
        if (!fs.existsSync(filePath))
          return { content: [{ type: "text", text: `File not found: ${filePath}` }], details: {} };

        await ensureDocOpen(filePath);

        const diags = lsp!.getDiagnostics(filePath);
        return {
          content: [{ type: "text", text: formatDiagnostics(diags, filePath) }],
          details: { count: diags.length },
        };
      }

      const text = formatAllDiagnostics(lsp!.getAllDiagnostics());
      return { content: [{ type: "text", text }], details: {} };
    },
  });

  // ── Tool: julia_complete ─────────────────────────────────────────────────

  pi.registerTool({
    name: "julia_complete",
    label: "Julia completions",
    description:
      "Get completion suggestions at a position in a Julia file. " +
      "Useful for discovering methods, fields, and exported names. " +
      "Requires the LSP server.",
    parameters: Type.Object({
      file: Type.String({ description: "Path to the .jl file." }),
      line: Type.Number({ description: "Line number (1-based)." }),
      col:  Type.Number({ description: "Column number (1-based)." }),
    }),
    async execute(_id, params, _signal, _onUpdate, ctx) {
      const g = lspGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      const filePath = path.resolve(ctx.cwd, params.file);
      if (!fs.existsSync(filePath))
        return { content: [{ type: "text", text: `File not found: ${filePath}` }], details: {} };

      await ensureDocOpen(filePath);
      const items = await lsp!.complete(filePath, params.line - 1, params.col - 1)
        .catch(() => []);

      const top = items.slice(0, 20).map((item) => {
        const detail = item.detail ? ` — ${item.detail}` : "";
        return `${item.label}${detail}`;
      });

      const text = top.length
        ? top.join("\n") + (items.length > 20 ? `\n…(${items.length - 20} more)` : "")
        : "No completions found.";

      return { content: [{ type: "text", text }], details: { total: items.length } };
    },
  });

  // ── Tool: julia_definition ───────────────────────────────────────────────

  pi.registerTool({
    name: "julia_definition",
    label: "Julia go-to-definition",
    description:
      "Find the definition of the symbol at a given position in a Julia file. " +
      "Returns the file path and line number. Requires the LSP server.",
    parameters: Type.Object({
      file: Type.String({ description: "Path to the .jl file." }),
      line: Type.Number({ description: "Line number (1-based)." }),
      col:  Type.Number({ description: "Column number (1-based)." }),
    }),
    async execute(_id, params, _signal, _onUpdate, ctx) {
      const g = lspGuard();
      if (g) return { content: [{ type: "text", text: g }], details: {} };

      const filePath = path.resolve(ctx.cwd, params.file);
      if (!fs.existsSync(filePath))
        return { content: [{ type: "text", text: `File not found: ${filePath}` }], details: {} };

      await ensureDocOpen(filePath);
      const locs = await lsp!.definition(filePath, params.line - 1, params.col - 1)
        .catch(() => []);

      if (locs.length === 0)
        return { content: [{ type: "text", text: "Definition not found." }], details: {} };

      const text = locs.map((loc) => {
        const file = loc.uri.replace("file://", "");
        const line = loc.range.start.line + 1;
        const col  = loc.range.start.character + 1;
        return `${file}:${line}:${col}`;
      }).join("\n");

      return { content: [{ type: "text", text }], details: { locations: locs.length } };
    },
  });
}