import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

interface ProviderConfig {
  id: string;
  baseUrl: string;
  envKey: string;
  suffix: string;
}

const PROVIDERS: ProviderConfig[] = [
  {
    id: "siliconflow",
    baseUrl: "https://api.siliconflow.com/v1",
    envKey: "SILICONFLOW_API_KEY",
    suffix: "SF"
  },
  {
    id: "deepinfra",
    baseUrl: "https://api.deepinfra.com/v1/openai",
    envKey: "DEEPINFRA_API_KEY",
    suffix: "DI"
  }
];

// Patterns matched in order; first match wins.
const CONTEXT_OVERRIDES: Array<[RegExp, number]> = [
  [/qwen3.*coder/i,           262144],
  [/qwen3.*next/i,            262144],
  [/qwen3-235b|qwen3-32b|qwen3-14b|qwen3-8b|qwen3-30b/i, 131072],
  [/qwq-32b/i,                131072],
  [/qwen2\.5-.*-128k/i,       131072],
  [/qwen2\.5/i,                32768],
  [/deepseek-v3|deepseek-r1/i, 131072],
  [/kimi-k2/i,                131072],
  [/glm-4\.[567]|glm-5/i,     131072],
  [/gpt-oss/i,                131072],
  [/minimax/i,                245760],
  [/hunyuan/i,                 32768],
];

function resolveContextWindow(modelId: string): number {
  for (const [pattern, ctx] of CONTEXT_OVERRIDES) {
    if (pattern.test(modelId)) return ctx;
  }
  return 32768; // conservative fallback
}

export default async function (pi: ExtensionAPI) {
  console.log("Provider discovery: Starting...");

  await Promise.all(PROVIDERS.map(async (config) => {
    const apiKey = process.env[config.envKey];
    if (!apiKey) {
      console.warn(`Provider discovery: ${config.envKey} not found. Skipping ${config.id}.`);
      return;
    }

    try {
      const response = await fetch(`${config.baseUrl}/models`, {
        headers: { "Authorization": `Bearer ${apiKey}` }
      });

      if (!response.ok) {
        throw new Error(`API error ${response.status}`);
      }

      const payload = await response.json();
      const models = payload.data
        .filter((model: any) => {
          // Heuristic to filter for chat/text models
          const id = model.id.toLowerCase();
          return id.includes("/") || model.type === "text-generation" || id.includes("gpt") || id.includes("llama") || id.includes("claude");
        })
        .map((model: any) => {
          const id = model.id;
          const isReasoning = id.toLowerCase().includes("r1") || 
                              id.toLowerCase().includes("thinking") || 
                              id.toLowerCase().includes("reasoning") ||
                              id.toLowerCase().includes("qwq");
          
          return {
            id: id,
            name: `${id.split('/').pop()} (${config.suffix})`,
            reasoning: isReasoning,
            input: ["text"],
            contextWindow: model.context_window ?? resolveContextWindow(id),
            maxTokens: 4096
          };
        });

      pi.registerProvider(config.id, {
        baseUrl: config.baseUrl,
        apiKey: config.envKey,
        api: "openai-completions",
        models: models,
        compat: {
          supportsDeveloperRole: false,
          supportsStrictMode: false,
          requiresToolResultName: true
        },
      });

      console.log(`Provider discovery: Registered ${models.length} models for ${config.id}.`);
    } catch (error) {
      console.error(`Provider discovery: Failed to fetch models for ${config.id}:`, error);
    }
  }));
}