#!/usr/bin/env node
'use strict';
const path = require('path');
const https = require('https');

const brain = require(path.resolve(__dirname, '..', '..', '..', 'brain'));

const LITELLM_URL = 'https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json';
const TIMEOUT_MS = 30000;

// Provider type mapping — LiteLLM litellm_provider values to our provider types
const PROVIDER_MAP = {
  openai: 'openai',
  anthropic: 'anthropic',
  'vertex_ai-anthropic_models': 'anthropic',
  'vertex_ai-language-models': 'google',
  gemini: 'google',
  'vertex_ai': 'google',
  ollama: 'ollama',
  ollama_chat: 'ollama',
};

/**
 * Fetch JSON from a URL via HTTPS. Returns parsed object.
 */
function fetchJson(url) {
  return new Promise((resolve, reject) => {
    const req = https.get(url, { timeout: TIMEOUT_MS }, (res) => {
      // Follow redirects
      if (res.statusCode >= 300 && res.statusCode < 400 && res.headers.location) {
        return fetchJson(res.headers.location).then(resolve, reject);
      }
      if (res.statusCode !== 200) {
        res.resume();
        return reject(new Error(`HTTP ${res.statusCode}`));
      }
      let data = '';
      res.on('data', chunk => { data += chunk; });
      res.on('end', () => {
        try { resolve(JSON.parse(data)); }
        catch (e) { reject(new Error('JSON parse error: ' + e.message)); }
      });
    });
    req.on('error', reject);
    req.on('timeout', () => { req.destroy(); reject(new Error('Request timeout')); });
  });
}

/**
 * Match a model_id from our registry against the LiteLLM catalog.
 * LiteLLM keys can be prefixed with provider (e.g., "openai/gpt-4o").
 * We try: exact match, provider-prefixed match, and partial match.
 */
function findInCatalog(catalog, modelId, providerType) {
  // 1. Exact match
  if (catalog[modelId]) return catalog[modelId];

  // 2. Provider-prefixed match (e.g., "openai/gpt-4o")
  const prefixes = Object.keys(PROVIDER_MAP).filter(k => PROVIDER_MAP[k] === providerType);
  for (const prefix of prefixes) {
    const key = `${prefix}/${modelId}`;
    if (catalog[key]) return catalog[key];
  }

  // 3. Partial match — find keys ending with our model_id
  for (const [key, entry] of Object.entries(catalog)) {
    if (key.endsWith('/' + modelId) || key === modelId) return entry;
  }

  return null;
}

// Models to skip — non-chat models, deprecated, or niche variants
const SKIP_PATTERNS = /^(text-|whisper|dall-e|tts-|babbage|davinci|canary|ft:|audio)/;
const SKIP_MODES = new Set(['embedding', 'image_generation', 'audio_transcription', 'audio_speech', 'moderation']);

/**
 * Infer capabilities from model name and LiteLLM entry.
 */
function inferCapabilities(modelId, entry) {
  const caps = ['chat'];
  if (entry.supports_function_calling || /gpt|claude|gemini|o[134]/.test(modelId)) caps.push('tools');
  if (entry.supports_vision || /vision|4o|4\.5|5|claude-3|gemini/.test(modelId)) caps.push('vision');
  if (/^o[134]|reasoning|thinking|claude-opus|claude-sonnet-4|pro/.test(modelId)) caps.push('reasoning');
  if (/codex|code|claude|gemini|gpt-4|gpt-5|o[134]/.test(modelId)) caps.push('code');
  return [...new Set(caps)];
}

/**
 * Infer speed tier (1-5) from model characteristics.
 */
function inferSpeedTier(modelId) {
  if (/mini|flash|haiku|nano|instant/.test(modelId)) return 5;
  if (/4o-mini|small|lite/.test(modelId)) return 4;
  if (/sonnet|4o(?!-mini)|flash/.test(modelId)) return 3;
  if (/pro|opus|5\.4|o3(?!-mini)/.test(modelId)) return 2;
  if (/ultra/.test(modelId)) return 1;
  return 3;
}

/**
 * Build a readable display name from a model ID.
 */
function buildDisplayName(modelId, providerType) {
  return modelId
    .replace(/^(openai|anthropic|google|gemini)\//, '')
    .replace(/^gpt-/, 'GPT-')
    .replace(/^o(\d)/, 'O$1')
    .replace(/^claude-/, 'Claude ')
    .replace(/^gemini-/, 'Gemini ')
    .replace(/-(\d{4})(\d{2})(\d{2})$/, '')  // strip date suffix
    .replace(/-/g, ' ')
    .replace(/\b\w/g, c => c.toUpperCase())
    .trim();
}

/**
 * Main sync function. Returns { updated, added, skipped, errors, total }.
 */
async function syncPricing() {
  console.log('[model-pricing-sync] Fetching LiteLLM pricing catalog...');
  const catalog = await fetchJson(LITELLM_URL);
  const catalogSize = Object.keys(catalog).length;
  console.log(`[model-pricing-sync] Catalog loaded: ${catalogSize} models`);

  const allModels = brain.listAllModels();
  const providers = brain.listModelProviders();
  let updated = 0;
  let added = 0;
  let skipped = 0;
  let errors = 0;

  // Phase 1: Update existing models
  for (const model of allModels) {
    const entry = findInCatalog(catalog, model.model_id, model.provider_type);
    if (!entry) {
      skipped++;
      continue;
    }

    const newInputCost = entry.input_cost_per_token ? entry.input_cost_per_token * 1_000_000 : null;
    const newOutputCost = entry.output_cost_per_token ? entry.output_cost_per_token * 1_000_000 : null;
    const newMaxContext = entry.max_input_tokens || entry.max_tokens || null;
    const newMaxOutput = entry.max_output_tokens || null;

    if (!newInputCost && !newOutputCost && !newMaxContext) {
      skipped++;
      continue;
    }

    const costChanged = newInputCost && Math.abs((model.cost_per_1m_input || 0) - newInputCost) > 0.001;
    const contextChanged = newMaxContext && model.max_context_tokens !== newMaxContext;
    const outputChanged = newMaxOutput && model.max_output_tokens !== newMaxOutput;

    if (!costChanged && !contextChanged && !outputChanged && model.cost_per_1m_input != null) {
      skipped++;
      continue;
    }

    try {
      let caps = model.capabilities;
      if (typeof caps === 'string') {
        try { caps = JSON.parse(caps); } catch { caps = []; }
      }

      brain.upsertModelRegistryEntry({
        id: model.id,
        providerId: model.provider_id,
        modelId: model.model_id,
        displayName: model.display_name,
        capabilities: caps || [],
        costPer1mInput: newInputCost || model.cost_per_1m_input,
        costPer1mOutput: newOutputCost || model.cost_per_1m_output,
        maxContextTokens: newMaxContext || model.max_context_tokens,
        maxOutputTokens: newMaxOutput || model.max_output_tokens,
        speedTier: model.speed_tier,
        enabled: model.enabled,
      });
      updated++;
      console.log(`[model-pricing-sync] Updated: ${model.display_name} — $${(newInputCost || 0).toFixed(2)}/$${(newOutputCost || 0).toFixed(2)} per 1M tokens`);
    } catch (err) {
      errors++;
      console.error(`[model-pricing-sync] Failed to update ${model.display_name}: ${err.message}`);
    }
  }

  // Phase 2: Discover new models from LiteLLM for existing providers
  const existingModelIds = new Set(allModels.map(m => m.model_id));
  const providerMap = {}; // type -> provider record
  for (const p of providers) {
    if (p.enabled) providerMap[p.type] = p;
  }

  for (const [catalogKey, entry] of Object.entries(catalog)) {
    if (!entry.litellm_provider) continue;
    if (entry.mode && SKIP_MODES.has(entry.mode)) continue;

    const ourType = PROVIDER_MAP[entry.litellm_provider];
    if (!ourType) continue;
    const provider = providerMap[ourType];
    if (!provider) continue;

    // Extract model ID from catalog key (strip provider prefix)
    const modelId = catalogKey.includes('/') ? catalogKey.split('/').slice(1).join('/') : catalogKey;
    if (SKIP_PATTERNS.test(modelId)) continue;
    if (existingModelIds.has(modelId)) continue;

    const inputCost = entry.input_cost_per_token ? entry.input_cost_per_token * 1_000_000 : null;
    const outputCost = entry.output_cost_per_token ? entry.output_cost_per_token * 1_000_000 : null;
    if (!inputCost && !outputCost) continue;

    // Only add significant chat models (skip obscure variants)
    const maxContext = entry.max_input_tokens || entry.max_tokens || null;
    if (maxContext && maxContext < 4000) continue;

    try {
      const displayName = buildDisplayName(modelId, ourType);
      brain.upsertModelRegistryEntry({
        id: `${provider.id}:${modelId}`,
        providerId: provider.id,
        modelId: modelId,
        displayName,
        capabilities: inferCapabilities(modelId, entry),
        costPer1mInput: inputCost,
        costPer1mOutput: outputCost,
        maxContextTokens: maxContext,
        maxOutputTokens: entry.max_output_tokens || null,
        speedTier: inferSpeedTier(modelId),
        enabled: 1,
      });
      existingModelIds.add(modelId);
      added++;
      console.log(`[model-pricing-sync] Added: ${displayName} (${ourType}) — $${(inputCost || 0).toFixed(2)}/$${(outputCost || 0).toFixed(2)}`);
    } catch (err) {
      // Likely a duplicate key — skip silently
      if (!err.message.includes('UNIQUE')) errors++;
    }
  }

  return { updated, added, skipped, errors, total: allModels.length, catalogSize };
}

// Entry point — called by skill-planner
if (require.main === module) {
  brain.initDb();
  syncPricing()
    .then(r => { console.log(`[model-pricing-sync] Done: ${r.updated} updated, ${r.added} added, ${r.skipped} skipped, ${r.errors} errors`); process.exit(0); })
    .catch(e => { console.error(`[model-pricing-sync] Fatal: ${e.message}`); process.exit(1); });
} else {
  module.exports = { run: syncPricing, syncPricing };
}
