// Provider entry + request hooks + catalog commands.
import { appendFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { NIM_API_KEY_REF, NIM_BASE_URL } from "./config/defaults";
import { applyCustomThinkingFormat, hasEnabledThinking } from "./handlers/thinking";
import type { TransformResult } from "./handlers/thinking";
import { findFamily } from "./config/model-families";
import { readCache, refreshCacheIfStale } from "./lib/metadata-cache";
import { createLogger } from "./lib/logger";
import {
	isCompactionSummarizationRequest,
	stripThinkingForSummarization,
} from "./lib/compaction-request.mjs";
import {
  STATIC_MODELS,
  STATIC_MODEL_MAP,
  classifyThinkingFormat,
} from "./models/registry";
import type { NimModelConfig } from "./models/types";
import {
	recordRetryAudit,
	setLastProviderRequestId,
} from "./extensions/lib/retry-audit.mjs";

const log = createLogger("nvidia-nim");
const NIM_DEBUG_LOG = join(homedir(), ".pi", "nim-debug.log");

const NVIDIA_NIM_PROVIDERS = new Set(["nvidia-nim", "nvidia"]);

function isNimProvider(ctx: ExtensionContext | undefined): boolean {
  const provider = ctx?.model?.provider;
  return Boolean(provider && NVIDIA_NIM_PROVIDERS.has(provider));
}

function injectExampleRequestExtra(
  payload: Record<string, unknown>,
  modelConfig: NimModelConfig,
  thinkingEnabled: boolean,
): boolean {
  const extra = modelConfig.exampleRequestExtra;
  if (!extra) return false;

  let modified = false;

  for (const [key, value] of Object.entries(extra)) {
    if (key === "chat_template_kwargs") {
      if (!thinkingEnabled) continue;
      const exampleKwargs = value as Record<string, unknown>;
      const kwargs = (payload.chat_template_kwargs as Record<string, unknown>) || {};
      let injected = false;
      for (const [kwKey, kwValue] of Object.entries(exampleKwargs)) {
        if (!(kwKey in kwargs)) {
          kwargs[kwKey] = kwValue;
          injected = true;
        }
      }
      if (injected) {
        payload.chat_template_kwargs = kwargs;
        modified = true;
      }
      continue;
    }

    if (!(key in payload)) {
      payload[key] = value;
      modified = true;
    }
  }

  return modified;
}

export function handleBeforeProviderRequest(event: { payload: unknown }) {
  const payload = event.payload as Record<string, unknown>;
  const modelId = payload.model as string | undefined;
  if (!modelId || !STATIC_MODEL_MAP.has(modelId)) return;

  normalizeContentArrays(payload);

  const modelConfig = STATIC_MODEL_MAP.get(modelId)!;

  const isSummarization = isCompactionSummarizationRequest(payload);
  if (isSummarization) {
    stripThinkingForSummarization(payload);
  }

  const format = classifyThinkingFormat(modelId);
  const result: TransformResult = applyCustomThinkingFormat(payload, format);
  let modified = result.modified;

  const thinkingEnabledAfterTransform = isSummarization
    ? false
    : result.thinkingEnabled !== undefined
      ? result.thinkingEnabled
      : hasEnabledThinking(payload);

  if (injectExampleRequestExtra(payload, modelConfig, thinkingEnabledAfterTransform)) {
    modified = true;
  }

  if (modelConfig.reasoningBudget != null && thinkingEnabledAfterTransform) {
    const budgetParamName = format === "thinking-budget" ? "thinking_budget" : "reasoning_budget";
    payload[budgetParamName] = modelConfig.reasoningBudget;
    modified = true;
  }

  if (payload.max_tokens == null && payload.max_completion_tokens == null) {
    payload.max_tokens = modelConfig.maxTokens;
    modified = true;
  }

  if (typeof process !== "undefined" && process.env.NIM_DEBUG) {
    try {
      appendFileSync(
        NIM_DEBUG_LOG,
        `--- ${new Date().toISOString()} ${modelId} ---\n` +
          JSON.stringify(payload, null, 2) + "\n",
      );
    } catch { /* ignore */ }
  }

  return modified ? payload : undefined;
}

export function handleAfterProviderResponse(
  event: { status: number; headers?: Record<string, string | undefined> },
  ctx: ExtensionContext,
): void {
  if (!isNimProvider(ctx)) return;

  const requestId =
    event.headers?.["x-request-id"] ?? event.headers?.["x-nvca-request-id"];
  if (requestId) {
    setLastProviderRequestId(requestId);
  }

  if (event.status === 429 || event.status >= 500) {
    recordRetryAudit("provider_response", {
      status: event.status,
      requestId,
    });
  }

  if (event.status === 429) {
    const retryAfter = event.headers?.["retry-after"];
    const notice = retryAfter
      ? `NVIDIA NIM rate-limited. Retry after ${retryAfter}.`
      : "NVIDIA NIM rate-limited.";
    ctx.ui.notify(notice, "warning");
    return;
  }

  if (event.status >= 500) {
    const notice = requestId
      ? `NVIDIA NIM server error (${event.status}). Request ID: ${requestId}`
      : `NVIDIA NIM server error (${event.status}).`;
    ctx.ui.notify(notice, "error");
  }
}

function normalizeContentArrays(payload: Record<string, unknown>): void {
  const messages = payload.messages as Array<Record<string, unknown>> | undefined;
  if (!messages) return;
  for (const msg of messages) {
    const content = msg.content;
    if (!Array.isArray(content)) continue;
    if (content.length === 0) {
      msg.content = "";
      continue;
    }
    const allText = content.every(
      (part) => (part as Record<string, unknown>).type === "text",
    );
    if (allText) {
      msg.content = content
        .map((part) => (part as Record<string, unknown>).text as string)
        .filter((text) => text != null)
        .join("\n");
    }
  }
}

function formatModelLine(model: NimModelConfig): string {
  const family = findFamily(model.id)?.name ?? "default";
  const format = classifyThinkingFormat(model.id);
  const thinking = model.reasoning ? "reasoning" : "plain";
  return `${model.id}  [${family}/${format}/${thinking}]`;
}

export default async function (pi: ExtensionAPI) {
  pi.registerProvider("nvidia-nim", {
    baseUrl: NIM_BASE_URL,
    apiKey: NIM_API_KEY_REF,
    api: "openai-completions",
    models: STATIC_MODELS,
  });

  pi.on("before_provider_request", (event) =>
    handleBeforeProviderRequest(event as { payload: unknown }),
  );
  pi.on("after_provider_response", (event, ctx) =>
    handleAfterProviderResponse(
      event as { status: number; headers?: Record<string, string | undefined> },
      ctx,
    ),
  );

  pi.on("session_start", async () => {
    const cache = await refreshCacheIfStale(false);
    if (cache) {
      const registered = STATIC_MODEL_MAP.size;
      const live = cache.models.length;
      log.info("catalog", `registered ${registered} models; live API lists ${live}`);
    }
  });

  pi.registerCommand("nim-refresh", {
    description: "Refresh NVIDIA NIM model catalog cache from the API",
    handler: async (_args, ctx) => {
      try {
        const cache = await refreshCacheIfStale(true);
        if (!cache) {
          ctx.ui.notify("Catalog refresh failed — check NVIDIA API key", "error");
          return;
        }
        ctx.ui.notify(
          `NIM catalog refreshed (${cache.models.length} models from API, ${STATIC_MODEL_MAP.size} registered)`,
          "info",
        );
      } catch (err) {
        ctx.ui.notify(`NIM refresh failed: ${err instanceof Error ? err.message : String(err)}`, "error");
      }
    },
  });

  pi.registerCommand("nim-models", {
    description: "List registered nvidia-nim models grouped by family",
    handler: async (_args, ctx) => {
      const byFamily = new Map<string, string[]>();
      for (const model of STATIC_MODELS) {
        const family = findFamily(model.id)?.name ?? "default";
        const lines = byFamily.get(family) ?? [];
        lines.push(formatModelLine(model));
        byFamily.set(family, lines);
      }

      const cache = readCacheSummary();
      const header = [
        `Registered: ${STATIC_MODELS.length} chat models`,
        cache ? `Live API (cached): ${cache.liveCount} models` : "Live API: not cached yet — run /nim-refresh",
        "",
        "Recommended free-dev stack:",
        "  nvidia-nim/nvidia-router                 (auto failover)",
        "  nvidia-nim/minimaxai/minimax-m3        (tools + vision)",
        "  nvidia-nim/nvidia/nemotron-3-super-120b-a12b  (reasoning)",
        "  nvidia-nim/meta/llama-3.3-70b-instruct (fast fallback)",
        "",
      ];

      const body: string[] = [];
      for (const [family, lines] of [...byFamily.entries()].sort(([a], [b]) => a.localeCompare(b))) {
        body.push(`## ${family} (${lines.length})`);
        body.push(...lines.sort());
        body.push("");
      }

      ctx.ui.notify(header.concat(body).join("\n"), "info");
    },
  });
}

function readCacheSummary(): { liveCount: number } | null {
  const cache = readCache();
  if (!cache) return null;
  return { liveCount: cache.models.length };
}