#!/usr/bin/env python3
"""
Script to map all :latest models to their underlying versions.
Usage:
  export CBORG_API_KEY=...
  python map_latest_models.py
"""
import os
import sys
from openai import OpenAI

def test_model_mapping(client, model_id):
    """Test a model and return the underlying model name."""
    try:
        response = client.chat.completions.create(
            model=model_id,
            messages=[{"role": "user", "content": "Hi"}],
            max_tokens=5
        )
        return response.model
    except Exception as e:
        return f"ERROR: {str(e)[:100]}"

def main():
    api_key = os.environ.get('CBORG_API_KEY')
    if not api_key:
        print("Error: CBORG_API_KEY environment variable not set.")
        sys.exit(1)

    client = OpenAI(
        api_key=api_key,
        base_url="https://api.cborg.lbl.gov"
    )
    
    # Get all available models
    try:
        response = client.models.list()
        all_models = [model.id for model in response.data]
    except Exception as e:
        print(f"Error fetching model list: {e}")
        sys.exit(1)
    
    # Filter for models with :latest
    latest_models = [m for m in all_models if ':latest' in m]
    
    # Also check models without suffix to compare
    base_models = []
    for latest in latest_models:
        base = latest.replace(':latest', '')
        if base in all_models:
            base_models.append(base)
    
    print("=" * 100)
    print("MAPPING OF :latest MODELS TO UNDERLYING VERSIONS")
    print("=" * 100)
    
    results = []
    
    # Test :latest models
    print(f"\nTesting {len(latest_models)} models with :latest suffix...")
    for model in sorted(latest_models):
        print(f"  Testing {model}...", end=" ", flush=True)
        underlying = test_model_mapping(client, model)
        results.append((model, underlying))
        print("✓")
    
    # Test base models for comparison
    print(f"\nTesting {len(base_models)} corresponding base models (without :latest)...")
    for model in sorted(base_models):
        print(f"  Testing {model}...", end=" ", flush=True)
        underlying = test_model_mapping(client, model)
        results.append((model, underlying))
        print("✓")
    
    # Print results
    print("\n" + "=" * 100)
    print("RESULTS")
    print("=" * 100)
    
    print("\n📋 Models with :latest suffix:")
    print("-" * 100)
    for model, underlying in results:
        if ':latest' in model:
            if underlying.startswith('ERROR'):
                print(f"❌ {model:<50} {underlying}")
            else:
                status = "→" if model != underlying else "="
                print(f"   {model:<50} {status} {underlying}")
    
    print("\n📋 Base models (without :latest):")
    print("-" * 100)
    for model, underlying in results:
        if ':latest' not in model:
            if underlying.startswith('ERROR'):
                print(f"❌ {model:<50} {underlying}")
            else:
                status = "→" if model != underlying else "="
                print(f"   {model:<50} {status} {underlying}")
    
    # Compare :latest vs base
    print("\n📊 COMPARISON: Do :latest and base versions map to the same model?")
    print("-" * 100)
    
    latest_map = {m: u for m, u in results if ':latest' in m}
    base_map = {m: u for m, u in results if ':latest' not in m}
    
    for latest, underlying_latest in sorted(latest_map.items()):
        base = latest.replace(':latest', '')
        if base in base_map:
            underlying_base = base_map[base]
            if underlying_latest == underlying_base:
                print(f"✓ {latest:<50} SAME as {base}")
                print(f"  └─ Both map to: {underlying_latest}")
            else:
                print(f"⚠️  {latest:<50} DIFFERENT from {base}")
                print(f"  ├─ :latest maps to: {underlying_latest}")
                print(f"  └─ base maps to:    {underlying_base}")
    
    print("\n" + "=" * 100)

if __name__ == '__main__':
    main()