import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import type { OAuthCredentials, OAuthLoginCallbacks } from "@mariozechner/pi-ai";

type NvidiaModelsResponse = {
  data?: Array<{
    id?: string;
    name?: string;
    context_window?: number;
    max_tokens?: number;
  }>;
};

const PROVIDER_NAME = "nvidia-build";
const DEFAULT_BASE_URL = "https://integrate.api.nvidia.com/v1";
const DEFAULT_CONTEXT_WINDOW = 128_000;
const DEFAULT_MAX_TOKENS = 16_384;

const FALLBACK_MODEL_IDS = [
  "meta/llama-3.1-70b-instruct",
  "meta/llama-3.3-70b-instruct",
  "nvidia/llama-3.1-nemotron-70b-instruct",
] as const;

const EXCLUDED_MODEL_HINTS = [
  "embedding",
  "embed",
  "rerank",
  "reward",
  "moderation",
  "safety",
  "tts",
  "asr",
  "whisper",
  "transcribe",
] as const;

function parseBoolean(value: string | undefined): boolean {
  if (!value) return false;
  return ["1", "true", "yes", "on"].includes(value.toLowerCase());
}

function splitCsv(value: string | undefined): string[] {
  if (!value) return [];
  return value
    .split(",")
    .map((part) => part.trim())
    .filter(Boolean);
}

function includeModel(id: string, includeAll: boolean): boolean {
  if (includeAll) return true;
  const lower = id.toLowerCase();
  return !EXCLUDED_MODEL_HINTS.some((hint) => lower.includes(hint));
}

function supportsVision(id: string): boolean {
  const lower = id.toLowerCase();
  return ["vision", "vl", "llava", "pixtral", "gemma-3", "qwen2-vl"].some((hint) => lower.includes(hint));
}

function supportsReasoning(id: string): boolean {
  const lower = id.toLowerCase();
  return ["reason", "thinking", "r1", "qwq"].some((hint) => lower.includes(hint));
}

function toDisplayName(id: string, name?: string): string {
  if (name && name.trim().length > 0) return name;
  return id
    .split("/")
    .pop()
    ?.replace(/[-_]+/g, " ")
    .replace(/\b\w/g, (char) => char.toUpperCase()) ?? id;
}

function toModelConfig(id: string, name?: string, contextWindow?: number, maxTokens?: number) {
  return {
    id,
    name: toDisplayName(id, name),
    reasoning: supportsReasoning(id),
    input: supportsVision(id) ? (["text", "image"] as const) : (["text"] as const),
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: contextWindow ?? DEFAULT_CONTEXT_WINDOW,
    maxTokens: maxTokens ?? DEFAULT_MAX_TOKENS,
    compat: {
      supportsDeveloperRole: false,
      supportsReasoningEffort: false,
      maxTokensField: "max_tokens" as const,
    },
  };
}

async function fetchModelsFromNvidia(baseUrl: string, apiKey: string, includeAll: boolean) {
  const response = await fetch(`${baseUrl}/models`, {
    headers: {
      Authorization: `Bearer ${apiKey}`,
      Accept: "application/json",
    },
  });

  if (!response.ok) {
    const body = await response.text();
    throw new Error(`NVIDIA models request failed (${response.status}): ${body}`);
  }

  const payload = (await response.json()) as NvidiaModelsResponse;
  const models = payload.data ?? [];

  return models
    .filter((model) => typeof model.id === "string" && model.id.length > 0)
    .filter((model) => includeModel(model.id!, includeAll))
    .map((model) => toModelConfig(model.id!, model.name, model.context_window, model.max_tokens));
}

function farFutureCredentials(key: string): OAuthCredentials {
  const tenYearsMs = 10 * 365 * 24 * 60 * 60 * 1000;
  return {
    refresh: key,
    access: key,
    expires: Date.now() + tenYearsMs,
  };
}

export default async function (pi: ExtensionAPI) {
  const baseUrl = (process.env.NVIDIA_BASE_URL || DEFAULT_BASE_URL).replace(/\/+$/, "");
  const includeAllModels = parseBoolean(process.env.NVIDIA_INCLUDE_ALL_MODELS);
  const csvModels = splitCsv(process.env.NVIDIA_MODELS);
  const envApiKey = process.env.NVIDIA_API_KEY || process.env.NVCF_API_KEY || "";
  const apiKeyEnvVar = process.env.NVIDIA_API_KEY ? "NVIDIA_API_KEY" : process.env.NVCF_API_KEY ? "NVCF_API_KEY" : "NVIDIA_API_KEY";

  const registerNvidiaProvider = (models: ReturnType<typeof toModelConfig>[]) => {
    pi.registerProvider(PROVIDER_NAME, {
      baseUrl,
      apiKey: apiKeyEnvVar,
      authHeader: true,
      api: "openai-completions",
      models,
      oauth: {
        name: "NVIDIA Build API",
        login: async (callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> => {
          const pasted = await callbacks.onPrompt({
            message: "Paste your NVIDIA API key:",
          });
          const apiKey = pasted.trim();
          if (!apiKey) throw new Error("No API key provided.");

          const discovered =
            csvModels.length > 0
              ? csvModels.map((id) => toModelConfig(id))
              : await fetchModelsFromNvidia(baseUrl, apiKey, includeAllModels);

          if (discovered.length === 0) {
            throw new Error("No models returned by NVIDIA for this API key.");
          }

          // Refresh provider models immediately after login so /model can show them.
          registerNvidiaProvider(discovered);

          return farFutureCredentials(apiKey);
        },
        refreshToken: async (credentials: OAuthCredentials): Promise<OAuthCredentials> => {
          const key = credentials.access || credentials.refresh;
          if (!key) throw new Error("Missing NVIDIA API key in stored credentials.");
          return farFutureCredentials(key);
        },
        getApiKey: (credentials: OAuthCredentials) => credentials.access,
      },
    });
  };

  let initialModels: ReturnType<typeof toModelConfig>[] = [];

  if (csvModels.length > 0) {
    initialModels = csvModels.map((id) => toModelConfig(id));
  } else if (envApiKey) {
    try {
      initialModels = await fetchModelsFromNvidia(baseUrl, envApiKey, includeAllModels);
    } catch {
      initialModels = FALLBACK_MODEL_IDS.map((id) => toModelConfig(id));
    }
  } else {
    initialModels = FALLBACK_MODEL_IDS.map((id) => toModelConfig(id));
  }

  registerNvidiaProvider(initialModels);
}