/**
 * The gpu_* tool surface.
 *
 * Tools are deliberately thin: they own the state machine (registry, ledger,
 * idempotency) and the provider command templates, and they enforce the spend
 * policy. All *how-to* knowledge — installing runtimes, downloading weights,
 * diagnosing CUDA failures — lives in the skills, not here.
 */
import { Type } from "typebox";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import {
  loadProviders, saveProviders, loadMachines, saveMachines, loadDeployments,
  saveDeployments, loadModels, loadWorkflows, appendJob, loadJobs,
  loadPolicy, savePolicy, monthSpendUsd, todaySpendUsd, nextJobId,
  nextDeploymentId, nextMachineId, fmtUsd, ProviderState, Deployment,
  Job, ModelManifest, WorkflowManifest,
} from "./registry.js";
import { PROVIDERS, runSpec, cliExists, parseInstanceId } from "./providers.js";

type Ctx = any; // ExtensionCommandContext — kept loose to stay robust across pi versions

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

/** Ask the user before anything that spends real money above the policy threshold. */
async function spendGate(ctx: Ctx, what: string, estUsd: number): Promise<boolean> {
  const policy = loadPolicy();
  if (estUsd <= policy.confirm_above_usd) return true;
  const ok = await ctx.ui.confirm("GPU spend", `${what} is estimated at ${fmtUsd(estUsd)}. Proceed?`, {
    detail: `Policy: confirm above ${fmtUsd(policy.confirm_above_usd)}.`,
    acceptLabel: "Proceed",
    rejectLabel: "Cancel",
  });
  return ok;
}

export function registerGpuTools(pi: ExtensionAPI) {
  // ------------------------------------------------------------------ setup
  pi.registerTool({
    name: "gpu_setup",
    label: "GPU Setup",
    description:
      "Detect installed GPU provider CLIs (runpodctl, vastai, modal, prime), check their authentication, and write the provider registry. Run this once per machine; re-run to rescan. If a CLI is missing, the skill provider-adapters tells you how to install it (ask permission first).",
    parameters: Type.Object({
      rescan: Type.Optional(Type.Boolean({ description: "Force re-detection of CLIs and auth (default false — uses cached results)" })),
    }),
    async execute(_id, params: { rescan?: boolean }, _sig, _upd, ctx: Ctx) {
      const providers = loadProviders();
      const out: string[] = ["GPU CLOUD SETUP", "-----------------"];
      const updated: ProviderState[] = [];
      const now = new Date().toISOString();

      for (const [pid, man] of Object.entries(PROVIDERS)) {
        const cached = providers.find((p) => p.id === pid);
        let installed = cached?.installed ?? false;
        let version: string | undefined = cached?.cliVersion;
        if (!cached || params.rescan) {
          if (man.cli) installed = await cliExists(man.cli);
        } else {
          installed = cached.installed;
        }
        let authenticated = cached?.authenticated ?? false;
        let authHint: string | undefined;
        if (installed && man.authCheck) {
          const r = await runSpec(man.authCheck, {}, 15000);
          authenticated = r.ok;
          if (!r.ok) authHint = (r.stderr || r.stdout || r.error || "").slice(0, 120);
        } else if (installed) {
          authenticated = true; // no auth check defined — treat CLI presence as enough
        }
        updated.push({ id: pid, name: man.name, kind: man.kind, cli: man.cli, installed, cliVersion: version, authenticated, authHint, detectedAt: now });
        const mark = installed ? "✓" : "○";
        const auth = authenticated ? "authenticated" : "NOT authenticated";
        out.push(`${mark} ${man.name.padEnd(10)} ${installed ? (man.cli || "").padEnd(10) + " " + auth : "CLI not installed"}`);
        if (authHint) out.push(`    hint: ${authHint}`);
      }
      saveProviders(updated);
      const ready = updated.filter((p) => p.installed && p.authenticated);
      out.push("-----------------");
      out.push(`Available providers: ${ready.length ? ready.map((p) => p.name).join(", ") : "none yet"}`);
      out.push("Install missing CLIs via the provider-adapters skill, then re-run /gpu setup.");
      return toolText(out.join("\n"), { providers: updated });
    },
  });

  // ----------------------------------------------------------------- status
  pi.registerTool({
    name: "gpu_status",
    label: "GPU Status",
    description:
      "Report the current GPU infrastructure: providers, deployments (with status), queued/running/completed jobs, and spend today and this month. Use before deciding to provision, reuse, or destroy anything.",
    parameters: Type.Object({}),
    async execute() {
      const providers = loadProviders();
      const deployments = loadDeployments();
      const jobs = loadJobs();
      const out: string[] = ["GPU INFRASTRUCTURE", "------------------"];
      out.push("Providers");
      for (const p of providers) out.push(`  ${p.name.padEnd(10)} ${p.installed && p.authenticated ? "✓ ready" : p.installed ? "○ installed, not authed" : "○ not installed"}`);
      out.push("Deployments");
      if (!deployments.length) out.push("  (none)");
      for (const d of deployments) {
        const dot = d.status === "ready" ? "●" : d.status === "provisioning" ? "◐" : d.status === "stopped" ? "○" : "✕";
        out.push(`  ${dot} ${d.id.padEnd(14)} ${d.model || d.runtime} on ${d.gpu} @ ${d.provider} ${d.endpoint ? d.endpoint : ""} (${fmtUsd(d.cost_hr)}/hr)`);
      }
      out.push(`Jobs  ${jobs.filter((j) => j.status === "running").length} running · ${jobs.filter((j) => j.status === "queued").length} queued · ${jobs.filter((j) => j.status === "completed").length} completed`);
      out.push(`Spend Today ${fmtUsd(todaySpendUsd())} · Month ${fmtUsd(monthSpendUsd())}`);
      const policy = loadPolicy();
      out.push(`Policy: per-job ≤ ${fmtUsd(policy.ceiling_per_job_usd)} · daily ≤ ${fmtUsd(policy.ceiling_daily_usd)} · monthly ≤ ${fmtUsd(policy.ceiling_monthly_usd)}`);
      return toolText(out.join("\n"), { deployments, providers });
    },
  });

  // -------------------------------------------------------------- provision
  pi.registerTool({
    name: "gpu_provision",
    label: "GPU Provision",
    description:
      "Provision a new GPU instance on a provider using that provider's CLI/API template (vastai create instance, runpodctl create pod, Lambda API). Records the machine and a provisioning deployment. Run gpu_setup first. For Vast, pass the bundle_id from a `vastai search offers` call; the provider-adapters skill shows the search command. Confirms with the user when the estimated hourly cost exceeds the spend policy.",
    parameters: Type.Object({
      provider: Type.String({ description: "Provider id: runpod | vast | lambda | prime (modal is serverless — see modal skill)", enum: ["runpod", "vast", "lambda", "prime"] }),
      gpu: Type.String({ description: "GPU type, e.g. RTX 5090, RTX 4090, H100 SXM, L40S" }),
      runtime: Type.String({ description: "Runtime to install: comfyui (default) or custom", default: "comfyui" }),
      model: Type.Optional(Type.String({ description: "Model id from models.json this instance is for (e.g. minimax-h3)" })),
      bundle_id: Type.Optional(Type.String({ description: "Vast.ai: bundle id from `vastai search offers`" })),
      template: Type.Optional(Type.String({ description: "RunPod: pod template name or id" })),
      disk_gb: Type.Optional(Type.Number({ description: "Disk size in GB (default 30)" })),
      region: Type.Optional(Type.String({ description: "Preferred region" })),
      cost_max_hr: Type.Optional(Type.Number({ description: "Refuse to provision above this hourly rate (USD)" })),
    }),
    async execute(_id, p: any, _sig, _upd, ctx: Ctx) {
      const man = PROVIDERS[p.provider];
      if (!man) return toolText(`Unknown provider: ${p.provider}`);
      if (man.kind === "serverless") return toolText("Modal is serverless — nothing to provision. See the modal skill: deploy is code (`modal deploy`), billing only while running.");
      if (!man.provision) return toolText(`${man.name}: no CLI provisioning template. See provider-adapters skill.`);

      const knownHr = estimateHourly(p.provider, p.gpu);
      if (p.cost_max_hr && knownHr > p.cost_max_hr) {
        return toolText(`Refusing: ${p.gpu} on ${man.name} ≈ ${fmtUsd(knownHr)}/hr exceeds cost_max_hr ${fmtUsd(p.cost_max_hr)}. Pick a cheaper GPU or raise the cap.`);
      }
      if (!(await spendGate(ctx, `Provisioning ${p.gpu} on ${man.name}`, knownHr * 2))) {
        return toolText("Cancelled by user.");
      }

      const vars: Record<string, string> = {
        gpu: p.gpu,
        template: p.template || "comfyui",
        disk: String(p.disk_gb || 30),
        bundle_id: p.bundle_id || "",
        region: p.region || "",
      };
      if (p.provider === "runpod") vars.cloud = "COMMUNITY";
      const r = await runSpec(man.provision, vars, 120000);
      const instanceId = parseInstanceId(p.provider, r.stdout);

      const machineId = nextMachineId();
      const machines = loadMachines();
      machines.push({
        id: machineId, provider: p.provider, instance_id: instanceId || "pending",
        gpu: p.gpu, gpu_count: 1, status: "running", cost_hr: knownHr,
        region: p.region, created_at: new Date().toISOString(), tags: [p.model || p.runtime],
      });
      saveMachines(machines);

      const depId = nextDeploymentId(p.model);
      const deployments = loadDeployments();
      deployments.push({
        id: depId, provider: p.provider, machine_id: machineId, instance_id: instanceId || undefined,
        gpu: p.gpu, gpu_count: 1, runtime: p.runtime, model: p.model,
        status: "provisioning", cost_hr: knownHr, region: p.region,
        created_at: new Date().toISOString(),
      });
      saveDeployments(deployments);

      const lines = [
        `PROVISION ${depId}`,
        `  provider: ${man.name}`,
        `  gpu:      ${p.gpu}`,
        `  instance: ${instanceId || "pending (check provider console)"}`,
        `  est cost: ${fmtUsd(knownHr)}/hr`,
      ];
      if (!r.ok) {
        lines.push(`  CLI error: ${(r.error || r.stderr).slice(0, 300)}`);
        lines.push("  Fall back to the provider-adapters skill (API flow) and register the instance with gpu_register.");
      } else {
        lines.push("  Next: wait for the instance to come up, install the runtime per the deployment/comfyui skills, then gpu_register with endpoint + ssh.");
      }
      return toolText(lines.join("\n"), { deployment_id: depId, machine_id: machineId, instance_id: instanceId, cli_ok: r.ok, cli_output: r.stdout.slice(0, 800) });
    },
  });

  // -------------------------------------------------------------- register
  pi.registerTool({
    name: "gpu_register",
    label: "GPU Register",
    description:
      "Register an existing GPU instance as a deployment (idempotent). Use this when the agent provisioned or configured an instance by hand/API and it is now ready: the deployment becomes the stable handle for gpu_run. Pass machine_id to re-register an instance created by gpu_provision that is now ready.",
    parameters: Type.Object({
      provider: Type.String({ description: "Provider id" }),
      instance_id: Type.Optional(Type.String({ description: "Provider's instance id (required if machine_id not given)" })),
      machine_id: Type.Optional(Type.String({ description: "Machine id from a previous gpu_provision" })),
      gpu: Type.String({ description: "GPU type, e.g. RTX 5090" }),
      runtime: Type.String({ description: "Runtime: comfyui | native | custom" }),
      model: Type.Optional(Type.String({ description: "Model id" })),
      endpoint: Type.Optional(Type.String({ description: "HTTP endpoint (ComfyUI: http://host:8188)" })),
      ssh: Type.Optional(Type.String({ description: "SSH target, e.g. root@1.2.3.4" })),
      cost_hr: Type.Optional(Type.Number({ description: "Actual hourly cost (USD)" })),
      status: Type.Optional(Type.String({ description: "ready (default) | provisioning | stopped", default: "ready" })),
    }),
    async execute(_id, p: any) {
      const deployments = loadDeployments();
      const existing = deployments.find((d) => d.instance_id === p.instance_id && d.instance_id !== "pending") ||
        deployments.find((d) => d.machine_id === p.machine_id);
      let dep: Deployment;
      if (existing) {
        dep = { ...existing, gpu: p.gpu || existing.gpu, runtime: p.runtime || existing.runtime, model: p.model || existing.model, endpoint: p.endpoint || existing.endpoint, ssh: p.ssh || existing.ssh, cost_hr: p.cost_hr ?? existing.cost_hr, status: (p.status || "ready") as Deployment["status"] };
        saveDeployments(deployments.map((d) => (d.id === existing.id ? dep : d)));
      } else {
        dep = {
          id: nextDeploymentId(p.model), provider: p.provider, instance_id: p.instance_id,
          gpu: p.gpu, gpu_count: 1, runtime: p.runtime, model: p.model,
          endpoint: p.endpoint, ssh: p.ssh, status: (p.status || "ready") as Deployment["status"],
          cost_hr: p.cost_hr ?? estimateHourly(p.provider, p.gpu),
          created_at: new Date().toISOString(),
        };
        deployments.push(dep);
        saveDeployments(deployments);
      }
      const lines = [
        `DEPLOYMENT ${dep.id}`,
        `  provider: ${dep.provider}`,
        `  gpu:      ${dep.gpu}`,
        `  runtime:  ${dep.runtime}`,
        `  model:    ${dep.model || "—"}`,
        `  endpoint: ${dep.endpoint || "—"}`,
        `  ssh:      ${dep.ssh || "—"}`,
        `  cost:     ${fmtUsd(dep.cost_hr)}/hr`,
        `  status:   ${dep.status.toUpperCase()}`,
      ];
      return toolText(lines.join("\n"), { deployment: dep });
    },
  });

  // --------------------------------------------------------------- destroy
  pi.registerTool({
    name: "gpu_destroy",
    label: "GPU Destroy",
    description:
      "Destroy a GPU machine or deployment (terminates billing). Confirms with the user if the deployment is still serving jobs or if monthly spend is high. After destruction the deployment stays in the registry with status destroyed — history and ledger are preserved.",
    parameters: Type.Object({
      id: Type.String({ description: "Deployment id (dep_...) or machine id (m_...)" }),
    }),
    async execute(_id, p: { id: string }, _sig, _upd, ctx: Ctx) {
      const deployments = loadDeployments();
      const machines = loadMachines();
      const dep = deployments.find((d) => d.id === p.id);
      const machine = machines.find((m) => m.id === p.id) || (dep ? machines.find((m) => m.id === dep.machine_id) : undefined);
      if (!dep && !machine) return toolText(`Nothing found for "${p.id}".`);

      const recentJobs = loadJobs().filter((j) => (dep && j.deployment_id === dep.id) && (j.status === "running" || j.status === "queued"));
      if (recentJobs.length) {
        const ok = await ctx.ui.confirm("GPU destroy", `${recentJobs.length} job(s) still running/queued on ${p.id}. Destroy anyway?`, { acceptLabel: "Destroy", rejectLabel: "Cancel" });
        if (!ok) return toolText("Cancelled.");
      }

      const man = dep ? PROVIDERS[dep.provider] : undefined;
      let cliResult = "";
      if (man?.destroy && machine?.instance_id && machine.instance_id !== "pending") {
        const r = await runSpec(man.destroy, { id: machine.instance_id }, 60000);
        cliResult = r.ok ? "destroy command ok" : `destroy command failed: ${(r.error || r.stderr).slice(0, 200)}`;
      }

      if (machine) saveMachines(machines.map((m) => (m.id === machine.id ? { ...m, status: "destroyed" as const } : m)));
      if (dep) saveDeployments(deployments.map((d) => (d.id === dep.id ? { ...d, status: "destroyed" as const, endpoint: undefined } : d)));
      return toolText(`Destroyed ${p.id}. ${cliResult}\nRecord kept (status destroyed). Jobs and cost history remain in the ledger.`);
    },
  });

  // ----------------------------------------------------------------- ensure
  pi.registerTool({
    name: "gpu_ensure",
    label: "GPU Ensure",
    description:
      "The idempotent entry point: given a model (and optionally workflow constraints), return a READY deployment that can run it — reusing an existing one or producing a provisioning plan (cheapest provider per the pricing docs, GPU class from the model manifest). Prefer this over gpu_provision when the goal is 'make model X runnable'. If create=true and a plan is produced, follow it: run the search/provision steps from the provider-adapters skill, then gpu_register.",
    parameters: Type.Object({
      model: Type.String({ description: "Model id from models.json (e.g. wan-2.2-5b, minimax-h3, hunyuan-video)" }),
      workflow: Type.Optional(Type.String({ description: "Workflow id to run (e.g. h3-image-to-video)" })),
      gpu: Type.Optional(Type.String({ description: "Force a GPU class (default: cheapest class that fits the model's VRAM)" })),
      provider: Type.Optional(Type.String({ description: "Force a provider" })),
      cost_max_hr: Type.Optional(Type.Number({ description: "Max hourly rate (USD)" })),
      create: Type.Optional(Type.Boolean({ description: "Produce a provisioning plan if nothing ready exists (default true)" })),
    }),
    async execute(_id, p: any) {
      const models = loadModels();
      const model = models.find((m) => m.id === p.model);
      if (!model) {
        return toolText(`Model "${p.model}" not in registry. Add it with gpu_model_add (or read the model-deployment skill to discover it first).`);
      }
      const deployments = loadDeployments();
      const ready = deployments.filter((d) => d.status === "ready" && (!p.workflow || d.model === p.model) && (!p.gpu || d.gpu === p.gpu));
      const reuse = ready.find((d) => !p.provider || d.provider === p.provider) || ready[0];
      if (reuse) {
        return toolText(
          [
            `REUSE ${reuse.id}`,
            `  ${reuse.gpu} on ${reuse.provider} — ${reuse.runtime}${reuse.model ? ` for ${reuse.model}` : ""}`,
            `  endpoint: ${reuse.endpoint || "—"}`,
            `  cost: ${fmtUsd(reuse.cost_hr)}/hr`,
          ].join("\n"),
          { deployment: reuse, action: "reuse" },
        );
      }
      if (p.create === false) return toolText(`No ready deployment for ${p.model}.`, { action: "none" });

      const fits = model.vram_gb_fp8 || model.vram_gb_min;
      const gpuSuggestion = p.gpu || defaultGpuForVram(fits);
      const providerSuggestion = p.provider || cheapestProviderFor(gpuSuggestion);
      const plan = [
        `PLAN: provision ${model.name} (${model.gpu_class}, ≥${model.vram_gb_min}GB VRAM)`,
        `  1. gpu_provision: provider=${providerSuggestion}, gpu=${gpuSuggestion} (or run the provider-adapters search for the cheapest live offer)`,
        `  2. Install runtime + weights per the deployment/comfyui skills (agent-driven)`,
        `  3. gpu_register with endpoint/ssh → becomes READY`,
        `  4. gpu_run with workflow ${p.workflow || "<workflow>"}`,
        `  Reference pricing: docs/pricing/${providerSuggestion}.md (snapshot — verify live before committing spend)`,
      ];
      if (providerSuggestion === "vast" && /4090|5090|3090/.test(gpuSuggestion)) {
        plan.splice(5, 0, `  Cheaper lane: Salad rents the same 24-32GB cards at ~2x less (batch pricing) — docs/pricing/salad.md. Aggressive option for retryable batches.`);
      }
      return toolText(plan.join("\n"), { action: "provision", plan: { provider: providerSuggestion, gpu: gpuSuggestion } });
    },
  });

  // -------------------------------------------------------------------- run
  pi.registerTool({
    name: "gpu_run",
    label: "GPU Run",
    description:
      "Submit a job to a deployment. Creates the ledger entry (J-...) with inputs and status running. Execution itself is agent-driven per the workflow skill (ComfyUI: POST /prompt, poll /history; native: HTTP call; provider API: provider call). When the run finishes, call gpu_job_finish with the output URIs so the ledger is complete.",
    parameters: Type.Object({
      deployment_id: Type.String({ description: "Deployment id (dep_...)" }),
      workflow: Type.String({ description: "Workflow id (e.g. h3-image-to-video)" }),
      inputs: Type.Object({}, { additionalProperties: true, description: "Workflow inputs: prompt, image, duration, etc." }),
      model: Type.Optional(Type.String({ description: "Model id (defaults to the deployment's model)" })),
      artifacts_expected: Type.Optional(Type.Array(Type.String({ description: "Output artifact names, e.g. output.mp4" }))),
    }),
    async execute(_id, p: any) {
      const deployments = loadDeployments();
      const dep = deployments.find((d) => d.id === p.deployment_id);
      if (!dep) return toolText(`Unknown deployment ${p.deployment_id}.`);
      if (dep.status !== "ready") return toolText(`Deployment ${dep.id} is ${dep.status}, not ready. Ensure it first.`);
      const now = new Date().toISOString();
      const job: Job = {
        id: nextJobId(),
        deployment_id: dep.id,
        provider: dep.provider,
        model: p.model || dep.model,
        workflow: p.workflow,
        inputs: p.inputs || {},
        status: "running",
        started_at: now,
        gpu: dep.gpu,
        instance_id: dep.instance_id,
        artifacts: [],
        ts: now,
      };
      appendJob(job);
      return toolText(
        [
          `JOB ${job.id}`,
          `  workflow:   ${job.workflow}`,
          `  deployment: ${dep.id} (${dep.gpu} @ ${dep.provider})`,
          `  started:    ${now}`,
          `  inputs:     ${JSON.stringify(job.inputs).slice(0, 200)}`,
          `Execute per the workflow skill, then gpu_job_finish id=${job.id} with artifact URIs.`,
        ].join("\n"),
        { job },
      );
    },
  });

  // ------------------------------------------------------------ job finish
  pi.registerTool({
    name: "gpu_job_finish",
    label: "GPU Job Finish",
    description:
      "Close out a job in the ledger: mark completed/failed, record duration, actual cost (or estimate from deployment hourly rate × duration), and output artifact URIs. Required for accurate spend rollups.",
    parameters: Type.Object({
      job_id: Type.String({ description: "Job id (J-...)" }),
      status: Type.String({ description: "completed | failed | cancelled", enum: ["completed", "failed", "cancelled"] }),
      artifacts: Type.Optional(Type.Array(Type.Object({ name: Type.String(), uri: Type.String() }), { description: "Output artifacts (durable URIs — copy off the pod before shutdown!)" })),
      duration_s: Type.Optional(Type.Number({ description: "Wall-clock duration in seconds" })),
      cost_usd: Type.Optional(Type.Number({ description: "Actual cost; omit to estimate from deployment rate" })),
      error: Type.Optional(Type.String({ description: "Failure reason" })),
      workflow_version: Type.Optional(Type.Number({ description: "Workflow version used (for 'regenerate everything from v3' queries)" })),
    }),
    async execute(_id, p: any) {
      const jobs = loadJobs();
      const idx = jobs.findIndex((j) => j.id === p.job_id);
      if (idx < 0) return toolText(`Unknown job ${p.job_id}.`);
      const job = jobs[idx];
      const dep = loadDeployments().find((d) => d.id === job.deployment_id);
      const duration = p.duration_s ?? ((Date.now() - new Date(job.started_at).getTime()) / 1000);
      const cost = p.cost_usd ?? (p.status === "completed" ? Math.max(0.001, (duration / 3600) * (dep?.cost_hr || 1)) : 0);
      const completed: Job = {
        ...job,
        status: p.status,
        completed_at: new Date().toISOString(),
        duration_s: Math.round(duration * 10) / 10,
        cost_usd: Math.round(cost * 1000) / 1000,
        artifacts: p.artifacts || job.artifacts,
        error: p.error,
        workflow_version: p.workflow_version ?? job.workflow_version,
      };
      // rewrite the JSONL line in place (keeps the ledger single-source)
      const fs = await import("node:fs");
      const path = `${process.env.GPU_STATE_DIR || `${process.cwd()}/.pi/gpu`}/jobs.jsonl`;
      const lines = fs.readFileSync(path, "utf-8").split("\n").filter(Boolean);
      lines[idx] = JSON.stringify(completed);
      fs.writeFileSync(path, lines.join("\n") + "\n");
      return toolText(
        [
          `JOB ${job.id} → ${p.status.toUpperCase()}`,
          `  duration: ${completed.duration_s}s`,
          `  cost:     ${fmtUsd(completed.cost_usd || 0)}`,
          `  artifacts:${completed.artifacts.length ? "\n    " + completed.artifacts.map((a) => `${a.name} → ${a.uri}`).join("\n    ") : " none"}`,
          p.error ? `  error:    ${p.error}` : "",
        ].filter(Boolean).join("\n"),
        { job: completed },
      );
    },
  });

  // ------------------------------------------------------------------- jobs
  pi.registerTool({
    name: "gpu_jobs",
    label: "GPU Jobs",
    description:
      "Query the job ledger: list jobs with filters (status, deployment, model, workflow), optionally summing cost. Use for 'how much did the last 100 videos cost', 'which jobs used workflow v3', 'what is queued'.",
    parameters: Type.Object({
      status: Type.Optional(Type.String({ description: "Filter: running | queued | completed | failed" })),
      deployment_id: Type.Optional(Type.String()),
      model: Type.Optional(Type.String()),
      workflow: Type.Optional(Type.String()),
      workflow_version: Type.Optional(Type.Number()),
      since: Type.Optional(Type.String({ description: "ISO date or 'today' | 'month'" })),
      limit: Type.Optional(Type.Number({ description: "Max rows (default 25)", default: 25 })),
      sum_cost: Type.Optional(Type.Boolean({ description: "Append total cost of the matched set" })),
    }),
    async execute(_id, p: any) {
      let jobs = loadJobs();
      if (p.status) jobs = jobs.filter((j) => j.status === p.status);
      if (p.deployment_id) jobs = jobs.filter((j) => j.deployment_id === p.deployment_id);
      if (p.model) jobs = jobs.filter((j) => j.model === p.model);
      if (p.workflow) jobs = jobs.filter((j) => j.workflow === p.workflow);
      if (p.workflow_version !== undefined) jobs = jobs.filter((j) => j.workflow_version === p.workflow_version);
      if (p.since === "today") jobs = jobs.filter((j) => j.ts.slice(0, 10) === new Date().toISOString().slice(0, 10));
      else if (p.since === "month") jobs = jobs.filter((j) => j.ts.slice(0, 7) === new Date().toISOString().slice(0, 7));
      else if (p.since) jobs = jobs.filter((j) => j.ts >= p.since);
      const total = jobs.reduce((s, j) => s + (j.cost_usd || 0), 0);
      const rows = jobs.slice(-(p.limit || 25)).reverse();
      const out = rows.map((j) =>
        `${j.id} ${j.status.padEnd(9)} ${(j.workflow || "").slice(0, 22).padEnd(22)} ${j.model || ""} ${j.duration_s ? j.duration_s + "s" : ""} ${j.cost_usd !== undefined ? fmtUsd(j.cost_usd) : ""}${j.artifacts.length ? " → " + j.artifacts[0].uri : ""}`
      );
      if (!rows.length) return toolText(`No jobs match.`);
      const lines = [...out];
      if (p.sum_cost) lines.push(`TOTAL (${jobs.length} jobs): ${fmtUsd(total)}`);
      return toolText(lines.join("\n"), { count: jobs.length, total_cost_usd: Math.round(total * 1000) / 1000 });
    },
  });

  // ------------------------------------------------------------------ spend
  pi.registerTool({
    name: "gpu_spend",
    label: "GPU Spend",
    description:
      "Cost rollups from the job ledger, grouped by provider, GPU, model, or workflow. Use for 'which GPU is cheapest for this workload', 'how much did we spend on H3', 'what does this batch cost'.",
    parameters: Type.Object({
      period: Type.String({ description: "today | month | all (default month)", default: "month" }),
      group_by: Type.String({ description: "provider | gpu | model | workflow | deployment (default none — total only)", default: "none" }),
    }),
    async execute(_id, p: any) {
      let jobs = loadJobs().filter((j) => j.status === "completed");
      const now = new Date().toISOString();
      if (p.period === "today") jobs = jobs.filter((j) => j.ts.slice(0, 10) === now.slice(0, 10));
      else if (p.period === "month") jobs = jobs.filter((j) => j.ts.slice(0, 7) === now.slice(0, 7));
      const total = jobs.reduce((s, j) => s + (j.cost_usd || 0), 0);
      const out = [`SPEND ${p.period.toUpperCase()}: ${fmtUsd(total)} across ${jobs.length} completed jobs`];
      if (p.group_by !== "none") {
        const groups = new Map<string, { n: number; cost: number }>();
        for (const j of jobs) {
          const key = String((j as any)[p.group_by] || "unknown");
          const g = groups.get(key) || { n: 0, cost: 0 };
          g.n += 1; g.cost += j.cost_usd || 0;
          groups.set(key, g);
        }
        for (const [k, g] of [...groups.entries()].sort((a, b) => b[1].cost - a[1].cost)) {
          out.push(`  ${k.padEnd(20)} ${g.n.toString().padStart(4)} jobs  ${fmtUsd(g.cost)}`);
        }
      }
      return toolText(out.join("\n"), { total_cost_usd: Math.round(total * 1000) / 1000, count: jobs.length });
    },
  });

  // ----------------------------------------------------------------- models
  pi.registerTool({
    name: "gpu_model_add",
    label: "GPU Model Add",
    description: "Add or update a model manifest in models.json (VRAM requirements, source, runtime). Use when discovering a model that isn't in the registry.",
    parameters: Type.Object({
      id: Type.String({ description: "Stable id, e.g. wan-2.2-5b" }),
      name: Type.String({ description: "Display name" }),
      source_url: Type.String({ description: "Where the weights live (HF/GitHub)" }),
      vram_gb_min: Type.Number({ description: "Minimum VRAM in GB" }),
      gpu_class: Type.String({ description: "Recommended GPU class, e.g. 24GB class (RTX 4090/5090)" }),
      runtime: Type.Optional(Type.String({ description: "comfyui | native | provider-api", default: "comfyui" })),
      vram_gb_fp8: Type.Optional(Type.Number({ description: "VRAM in GB with fp8 quantization, if known" })),
      notes: Type.Optional(Type.String()),
      verified: Type.Optional(Type.Boolean({ description: "True if a test generation ran successfully", default: false })),
    }),
    async execute(_id, p: any) {
      const models = loadModels();
      const man: ModelManifest = { id: p.id, name: p.name, source_url: p.source_url, vram_gb_min: p.vram_gb_min, vram_gb_fp8: p.vram_gb_fp8, gpu_class: p.gpu_class, runtime: p.runtime || "comfyui", notes: p.notes || "", verified: p.verified || false };
      const i = models.findIndex((m) => m.id === p.id);
      if (i >= 0) models[i] = man; else models.push(man);
      const fs = await import("node:fs");
      const pth = `${process.env.GPU_STATE_DIR || `${process.cwd()}/.pi/gpu`}/models.json`;
      fs.writeFileSync(pth, JSON.stringify(models, null, 2) + "\n");
      return toolText(`Model ${p.id} saved (${p.name}, ≥${p.vram_gb_min}GB VRAM).`);
    },
  });

  // -------------------------------------------------------------- workflows
  pi.registerTool({
    name: "gpu_workflow_add",
    label: "GPU Workflow Add",
    description: "Register an executable workflow (ComfyUI JSON, HTTP contract, or provider API) so gpu_run can target it by id.",
    parameters: Type.Object({
      id: Type.String({ description: "Stable id, e.g. wan-i2v" }),
      runtime: Type.String({ description: "comfyui | http | provider-api | native" }),
      description: Type.String({ description: "What it does" }),
      model: Type.Optional(Type.String({ description: "Model id it belongs to" })),
      input_schema: Type.Optional(Type.Object({}, { additionalProperties: true, description: "Input field name → type" })),
    }),
    async execute(_id, p: any) {
      const workflows = loadWorkflows();
      const w: WorkflowManifest = { id: p.id, runtime: p.runtime, description: p.description, model: p.model, input_schema: p.input_schema || {}, verified: false };
      const i = workflows.findIndex((x) => x.id === p.id);
      if (i >= 0) workflows[i] = w; else workflows.push(w);
      const fs = await import("node:fs");
      const pth = `${process.env.GPU_STATE_DIR || `${process.cwd()}/.pi/gpu`}/workflows.json`;
      fs.writeFileSync(pth, JSON.stringify(workflows, null, 2) + "\n");
      return toolText(`Workflow ${p.id} saved (runtime: ${p.runtime}).`);
    },
  });

  // ----------------------------------------------------------------- policy
  pi.registerTool({
    name: "gpu_policy",
    label: "GPU Policy",
    description:
      "Read or update the spend policy: per-job ceiling, daily ceiling, monthly ceiling, and the threshold above which provisioning/destruction requires user confirmation. The agent cannot silently override these — they are the trust boundary for real-money actions.",
    parameters: Type.Object({
      ceiling_per_job_usd: Type.Optional(Type.Number()),
      ceiling_daily_usd: Type.Optional(Type.Number()),
      ceiling_monthly_usd: Type.Optional(Type.Number()),
      confirm_above_usd: Type.Optional(Type.Number()),
      idle_shutdown_after_min: Type.Optional(Type.Number({ description: "Minutes of idle before the agent should shut the machine down" })),
    }),
    async execute(_id, p: any) {
      const policy = loadPolicy();
      const updated = { ...policy, ...Object.fromEntries(Object.entries(p).filter(([, v]) => v !== undefined)) };
      savePolicy(updated);
      return toolText(
        [
          "SPEND POLICY",
          `  per-job ceiling:  ${fmtUsd(updated.ceiling_per_job_usd)}`,
          `  daily ceiling:    ${fmtUsd(updated.ceiling_daily_usd)}`,
          `  monthly ceiling:  ${fmtUsd(updated.ceiling_monthly_usd)}`,
          `  confirm above:    ${fmtUsd(updated.confirm_above_usd)}`,
          `  idle shutdown:    after ${updated.idle_shutdown_after_min} min`,
        ].join("\n"),
        { policy: updated },
      );
    },
  });
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function estimateHourly(provider: string, gpu: string): number {
  // Rough reference rates from docs/pricing snapshots (2026-08-12). The agent
  // should verify live rates at deploy time — these only feed confirmation gates.
  const refs: Record<string, Record<string, number>> = {
    runpod: { "RTX 5090": 0.99, "RTX 4090": 0.69, "RTX 3090": 0.46, "L40S": 0.99, "RTX A6000": 0.49, "H100 SXM": 2.99, "H100 PCIe": 2.89, "A100 SXM": 1.49, "H200": 4.39, "B200": 5.89, "RTX Pro 6000": 1.99, "RTX PRO 4500": 1.15, "L4": 0.44 },
    vast: { "RTX 5090": 0.31, "RTX 4090": 0.25, "RTX 3090": 0.08, "L40S": 0.43, "RTX A6000": 0.3, "H100 SXM": 1.0, "A100 SXM": 0.8, "H200 SXM": 1.2 },
    salad: { "RTX 5090": 0.294, "RTX 4090": 0.204, "RTX 3090": 0.124, "RTX PRO 6000": 0.8 },
    lambda: { "H100 SXM": 4.09, "B200 SXM6": 6.79, "A100 SXM": 2.79, "A6000": 1.09, "GH200": 2.29 },
    prime: { "H100": 2.43, "H200": 1.99, "B200": 3.49, "B300": 4.99 },
  };
  const g = gpu.toUpperCase();
  const row = refs[provider] || {};
  const exact = Object.entries(row).find(([k]) => g.includes(k.toUpperCase()));
  if (exact) return exact[1];
  return g.includes("5090") || g.includes("H100") || g.includes("B200") || g.includes("H200") ? 1.5 : 0.5;
}

function defaultGpuForVram(vram: number): string {
  if (vram >= 60) return "H100 SXM";
  if (vram >= 40) return "L40S";
  if (vram >= 24) return "RTX 5090";
  return "RTX 4090";
}

function cheapestProviderFor(gpu: string): string {
  const g = gpu.toUpperCase();
  if (g.includes("H100") || g.includes("H200") || g.includes("B200") || g.includes("B300")) return "vast";
  if (g.includes("5090") || g.includes("4090") || g.includes("L40S") || g.includes("A6000") || g.includes("3090")) return "vast";
  return "runpod";
}