import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { execSync } from "child_process";
import { readFileSync, existsSync } from "fs";
import { resolve } from "path";
import { homedir } from "os";

/**
 * Load configuration from ~/.pi/agent/settings.json.
 *
 * Looks for:
 *   { "litellm-autodiscover": { "baseUrl": "https://your-instance.example.com" } }
 *
 * Falls back to the built-in default if the key is absent or the file is unparsable.
 */
function loadConfig(): { baseUrl: string } {
  const settingsPath = resolve(homedir(), ".pi/agent/settings.json");
  if (existsSync(settingsPath)) {
    try {
      const settings = JSON.parse(readFileSync(settingsPath, "utf-8")) as Record<string, unknown>;
      const extConfig = settings["litellm-autodiscover"] as Record<string, unknown> | undefined;
      if (extConfig && typeof extConfig.baseUrl === "string" && extConfig.baseUrl.length > 0) {
        return { baseUrl: extConfig.baseUrl.replace(/\/+$/, "") };
      }
    } catch {
      console.warn("[LiteLLM] Could not parse settings.json, using built-in default baseUrl.");
    }
  }
  return { baseUrl: "https://llm.mlcloud.uni-tuebingen.de" };
}

const CONFIG = loadConfig();
const BASE_URL = CONFIG.baseUrl;
const MODELS_URL = `${BASE_URL}/v1/models`;
const MODEL_INFO_URL = `${BASE_URL}/v1/model/info`;

/**
 * Fetch model metadata from LiteLLM's /v1/model/info endpoint.
 * Returns a map of model_name -> { max_input_tokens, max_output_tokens, input_cost_per_token, output_cost_per_token, ... }
 */
async function fetchModelInfo(): Promise<Record<string, {
  max_input_tokens?: number;
  max_output_tokens?: number;
  input_cost_per_token?: number | null;
  output_cost_per_token?: number | null;
  cache_read_input_token_cost?: number | null;
  cache_creation_input_token_cost?: number | null;
  supports_vision?: boolean;
}>> {
  try {
    const apiKey = execSync("security find-generic-password -ws 'litellm'", { encoding: "utf-8" }).trim();
    const response = await fetch(MODEL_INFO_URL, {
      headers: { "Authorization": `Bearer ${apiKey}` }
    });

    if (!response.ok) {
      console.warn(`[LiteLLM] Failed to fetch model info (${response.status}), using defaults.`);
      return {};
    }

    const data = await response.json() as { data: Array<{ model_name: string; model_info?: Record<string, unknown> }> };
    const result: Record<string, {
      max_input_tokens?: number;
      max_output_tokens?: number;
      input_cost_per_token?: number | null;
      output_cost_per_token?: number | null;
      cache_read_input_token_cost?: number | null;
      cache_creation_input_token_cost?: number | null;
      supports_vision?: boolean;
    }> = {};

    for (const entry of data.data ?? []) {
      const modelInfo = entry.model_info as Record<string, unknown> | undefined;
      if (modelInfo) {
        const costData = {
          max_input_tokens: modelInfo.max_input_tokens as number | undefined,
          max_output_tokens: modelInfo.max_output_tokens as number | undefined,
          input_cost_per_token: (modelInfo.input_cost_per_token as number | null | undefined) ?? null,
          output_cost_per_token: (modelInfo.output_cost_per_token as number | null | undefined) ?? null,
          cache_read_input_token_cost: (modelInfo.cache_read_input_token_cost as number | null | undefined) ?? null,
          cache_creation_input_token_cost: (modelInfo.cache_creation_input_token_cost as number | null | undefined) ?? null,
          supports_vision: modelInfo.supports_vision as boolean | undefined,
        };
        // Index by model_name (e.g., "Qwen/Qwen3.6-35B-A3B")
        result[entry.model_name] = costData;
        // Also index by model_info.id (the hash) so the /v1/models list can look it up
        const infoId = modelInfo.id as string | undefined;
        if (infoId) {
          result[infoId] = costData;
        }
      }
    }

    return result;
  } catch {
    return {};
  }
}

export default async function (pi: ExtensionAPI) {
  try {
    // 1. Fetch available models from /v1/models
    const apiKey = execSync("security find-generic-password -ws 'litellm'", { encoding: "utf-8" }).trim();
    const modelsResponse = await fetch(MODELS_URL, {
      headers: { "Authorization": `Bearer ${apiKey}` }
    });

    if (!modelsResponse.ok) {
      console.error(`[LiteLLM] Failed to fetch models: ${modelsResponse.statusText}`);
      return;
    }

    const modelsData = await modelsResponse.json() as { data: Array<{ id: string; owned_by?: string }> };

    // 2. Fetch detailed model info from /v1/model/info (includes max_tokens, costs, etc.)
    console.log(`[LiteLLM] Fetched ${modelsData.data.length} model(s) from /v1/models.`);
    const modelInfo = await fetchModelInfo();

    // 3. Register the provider with discovered models
    pi.registerProvider("litellm", {
      name: "LiteLLM",
      baseUrl: `${BASE_URL}/v1`,
      apiKey: "!security find-generic-password -ws 'litellm'",
      api: "openai-completions",
      models: modelsData.data.map(model => {
        const isHash = /^[a-f0-9]{64}$/i.test(model.id);
        const displayName = isHash ? "LiteLLM Custom Model" : model.id;

        // Look up detailed info from the model/info endpoint
        const info = modelInfo[model.id] ?? {};
        const maxInput = info.max_input_tokens ?? 128000;
        const maxOutput = info.max_output_tokens ?? 8192;

        return {
          id: model.id,
          name: displayName,
          reasoning: false,
          input: info.supports_vision ? ["text", "image"] : ["text"],
          contextWindow: maxInput,
          maxTokens: maxOutput,
          cost: {
            // LiteLLM returns $/token, pi-ai's calculateCost divides by 1M,
            // so we convert to $/1M tokens
            input: (info.input_cost_per_token ?? 0) * 1_000_000,
            output: (info.output_cost_per_token ?? 0) * 1_000_000,
            cacheRead: (info.cache_read_input_token_cost ?? 0) * 1_000_000,
            cacheWrite: (info.cache_creation_input_token_cost ?? 0) * 1_000_000,
          }
        };
      })
    });

    console.log(`[LiteLLM] Successfully auto-discovered and registered ${modelsData.data.length} model(s).`);
  } catch (error) {
    console.error("[LiteLLM] Error setting up auto-discovery:", error instanceof Error ? error.message : error);
  }
}