#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 19 14:36:46 2025

@author: samd
"""
import bilbystats as bs
import json
from openai import OpenAI

import requests

def get_groq_models():
    # Replace with your actual API key
    API_KEY = bs.read_api_key('groq')

    url = "https://api.groq.com/openai/v1/models"
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    response = requests.get(url, headers=headers)

    groq_models = []
    
    if response.status_code == 200:
        models = response.json()
        for model in models.get("data", []):
            groq_models.append(model["id"])
    else:
        print("Error:", response.status_code, response.text)
        
    return groq_models

def get_openrouter_models(save_to_file: bool = False):
    client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=bs.read_api_key('openrouter'))
    response = client.models.list()

    # Access the data attribute which holds a list of model objects
    models = response.data

    # Print just the model IDs
    model_list = []
    for model in models:
        model_list.append(model.id)
        
    model_list = [model.id for model in response.data]

    # Save to JSON file
    if save_to_file:
        with open('/Users/samd/Documents/Packages/bilbystats/bilbystats/data/openrouter_models.json', 'w') as f:
            json.dump(model_list, f, indent=2)
        
    return model_list

def get_claude_models(use_api: bool = True) -> dict:
    """
    Get a dictionary of available Claude models for use with claude_api.
    
    This function returns information about Claude models that can be used
    with the claude_api function in bilbystats. It can either fetch the 
    current models from the Anthropic API or use a fallback list.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    use_api : bool, optional (default=True)
        Whether to fetch models from the Anthropic API. If False or if API
        call fails, uses a fallback list of known models.
    
    ---------------------------------------------------------------------------
    OUTPUT:
    models : dict
        A dictionary with the following structure:
        {
            'aliases': dict - Model aliases that claude_api accepts
            'available_models': list - Full list of Claude model identifiers
            'recommended': dict - Recommended models for different use cases
            'source': str - Whether models came from 'api' or 'fallback'
        }
    
    ---------------------------------------------------------------------------
    AUTHORS: Your Name
    ---------------------------------------------------------------------------
    """
    import bilbystats as bs
    import anthropic
    
    # Model aliases that claude_api currently accepts
    aliases = {
        "claude": "claude-3-7-sonnet-20250219",  # Default alias from your code
    }
    
    available_models = []
    
    try:
            # Try to get models from the Anthropic API
            key = bs.read_api_key("claude")
            client = anthropic.Anthropic(api_key=key)
            
            # Check if the client has a models list method
            if hasattr(client, 'models') and hasattr(client.models, 'list'):
                models_response = client.models.list()
                available_models = [model.id for model in models_response.data]
                source = "api"
                print(f"✅ Fetched {len(available_models)} models from Anthropic API")
            else:
                print("⚠️  Anthropic client doesn't support model listing, using fallback")
                use_api = False
                
    except Exception as e:
            print(f"⚠️  Could not fetch models from API: {e}")
            use_api = False
  
    # Try to identify recommended models from the available list
    recommended = {}
    
    # Look for specific patterns in available models
    sonnet_models = [m for m in available_models if 'sonnet' in m.lower()]
    haiku_models = [m for m in available_models if 'haiku' in m.lower()]
    opus_models = [m for m in available_models if 'opus' in m.lower()]
    
    if sonnet_models:
        # Get the latest sonnet (usually has the most recent date)
        latest_sonnet = sorted(sonnet_models, reverse=True)[0]
        recommended['latest_sonnet'] = latest_sonnet
    
    if opus_models:
        recommended['most_capable'] = sorted(opus_models, reverse=True)[0]
    
    if haiku_models:
        latest_haiku = sorted(haiku_models, reverse=True)[0]
        recommended['fastest'] = latest_haiku
    
    # Always include the current default
    recommended['current_default'] = "claude-3-7-sonnet-20250219"
    
    return {
        'aliases': aliases,
        'available_models': available_models,
        'recommended': recommended,
        'source': source
    }

def get_openai_models() -> dict:
    """
    Get a dictionary of available OpenAI models for use with openai_api.
    
    This function returns information about OpenAI models that can be used
    with the openai_api function in bilbystats. It fetches the current models
    from the OpenAI API.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    use_api : bool, optional (default=True)
        Whether to fetch models from the OpenAI API. If False, returns empty list.
    
    ---------------------------------------------------------------------------
    OUTPUT:
    models : dict
        A dictionary with the following structure:
        {
            'aliases': dict - Model aliases that openai_api accepts
            'available_models': list - Full list of OpenAI model identifiers
            'recommended': dict - Recommended models for different use cases
            'source': str - Whether models came from 'api' or 'unavailable'
        }
    
    ---------------------------------------------------------------------------
    AUTHORS: Your Name
    ---------------------------------------------------------------------------
    """
    import bilbystats as bs
    import openai
    
    available_models = []
    source = "unavailable"
    
    
    try:
            # Try to get models from the OpenAI API
            key = bs.read_api_key("openai")
            client = openai.OpenAI(api_key=key)
            
            # Get models from OpenAI API
            models_response = client.models.list()
            available_models = [model.id for model in models_response.data]
            source = "api"
            print(f"✅ Fetched {len(available_models)} models from OpenAI API")
                
    except Exception as e:
            print(f"⚠️  Could not fetch models from OpenAI API: {e}")
            available_models = []
            source = "error"
    
    # Try to identify recommended models from the available list
    recommended = {}
    
    if available_models:
        # Look for specific patterns in available models
        gpt4_models = [m for m in available_models if 'gpt-4' in m.lower()]
        gpt35_models = [m for m in available_models if 'gpt-3.5' in m.lower()]
        gpt4o_models = [m for m in available_models if 'gpt-4o' in m.lower()]
        
        # Prioritize GPT-4o models
        if gpt4o_models:
            # Look for specific variants
            gpt4o_mini = [m for m in gpt4o_models if 'mini' in m.lower()]
            gpt4o_regular = [m for m in gpt4o_models if 'mini' not in m.lower()]
            
            if gpt4o_regular:
                recommended['most_capable'] = sorted(gpt4o_regular)[-1]  # Latest
                recommended['balanced'] = sorted(gpt4o_regular)[-1]
            if gpt4o_mini:
                recommended['cost_effective'] = sorted(gpt4o_mini)[-1]
                recommended['fastest'] = sorted(gpt4o_mini)[-1]
        
        # GPT-4 models as backup
        elif gpt4_models:
            recommended['most_capable'] = sorted(gpt4_models)[-1]
            recommended['balanced'] = sorted(gpt4_models)[-1]
        
        # GPT-3.5 as cost-effective option
        if gpt35_models:
            recommended['cost_effective'] = sorted(gpt35_models)[-1]
        
        # Set current defaults based on your code
        if 'gpt-4o' in available_models:
            recommended['current_default'] = 'gpt-4o'
        elif 'gpt-4o-mini' in available_models:
            recommended['current_default'] = 'gpt-4o-mini'
            
        
        recommended['current_default'] = 'gpt-4o-mini'
        recommended['most_capable'] = 'gpt-4o'
    
    return {
        'available_models': available_models,
        'recommended': recommended,
        'source': source
    }

def get_gemini_models(use_api: bool = True) -> dict:
    """
    Get a dictionary of available Gemini models for use with gemini_api.
    
    This function returns information about Gemini models that can be used
    with the gemini_api function in bilbystats. It fetches the current models
    from the Google Gemini API.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    use_api : bool, optional (default=True)
        Whether to fetch models from the Gemini API. If False, returns empty list.
    
    ---------------------------------------------------------------------------
    OUTPUT:
    models : dict
        A dictionary with the following structure:
        {
            'aliases': dict - Model aliases that gemini_api accepts
            'available_models': list - Full list of Gemini model identifiers
            'recommended': dict - Recommended models for different use cases
            'source': str - Whether models came from 'api' or 'unavailable'
        }
    
    ---------------------------------------------------------------------------
    AUTHORS: Your Name
    ---------------------------------------------------------------------------
    """
    import bilbystats as bs
    from google import genai
    
    # Model aliases that gemini_api currently accepts
    aliases = {
        "gemini": "gemini-2.0-flash",  # From your code
    }
    
    available_models = []
    source = "unavailable"
    
    if use_api:
        try:
            # Try to get models from the Gemini API
            key = bs.read_api_key("gemini")
            client = genai.Client(api_key=key)
            
            # Get models from Gemini API - handle Pager object
            models_response = client.models.list()
            
            # Handle different response types
            if hasattr(models_response, 'models'):
                # Direct models attribute
                model_list = models_response.models
            elif hasattr(models_response, '__iter__'):
                # Pager object - iterate through it
                model_list = list(models_response)
            else:
                # Try to access as list directly
                model_list = models_response
            
            # Extract model names, handling different name formats
            available_models = []
            for model in model_list:
                if hasattr(model, 'name'):
                    # Extract just the model name from full path (e.g., "models/gemini-pro" -> "gemini-pro")
                    model_name = model.name.split('/')[-1] if '/' in model.name else model.name
                    available_models.append(model_name)
                elif hasattr(model, 'id'):
                    available_models.append(model.id)
                else:
                    # Fallback - convert to string
                    available_models.append(str(model))
            
            source = "api"
            print(f"✅ Fetched {len(available_models)} models from Gemini API")
                
        except Exception as e:
            print(f"⚠️  Could not fetch models from Gemini API: {e}")
            available_models = []
            source = "error"
    
    # Try to identify recommended models from the available list
    recommended = {}
    
    if available_models:
        # Look for specific patterns in available models
        gemini2_models = [m for m in available_models if 'gemini-2' in m.lower()]
        gemini15_models = [m for m in available_models if 'gemini-1.5' in m.lower()]
        flash_models = [m for m in available_models if 'flash' in m.lower()]
        pro_models = [m for m in available_models if 'pro' in m.lower()]
        
        # Prioritize Gemini 2.0 models
        if gemini2_models:
            flash_2_models = [m for m in gemini2_models if 'flash' in m.lower()]
            pro_2_models = [m for m in gemini2_models if 'pro' in m.lower()]
            
            if flash_2_models:
                recommended['balanced'] = sorted(flash_2_models)[-1]
                recommended['fastest'] = sorted(flash_2_models)[-1]
            if pro_2_models:
                recommended['most_capable'] = sorted(pro_2_models)[-1]
        
        # Gemini 1.5 as backup
        elif gemini15_models:
            flash_15_models = [m for m in gemini15_models if 'flash' in m.lower()]
            pro_15_models = [m for m in gemini15_models if 'pro' in m.lower()]
            
            if pro_15_models:
                recommended['most_capable'] = sorted(pro_15_models)[-1]
            if flash_15_models:
                recommended['balanced'] = sorted(flash_15_models)[-1]
                recommended['fastest'] = sorted(flash_15_models)[-1]
        
        # General patterns
        if flash_models:
            recommended['cost_effective'] = sorted(flash_models)[-1]
        
        # Set current default based on your code
        if 'gemini-2.0-flash' in available_models:
            recommended['current_default'] = 'gemini-2.0-flash'
        elif any('gemini-2' in m for m in available_models):
            recommended['current_default'] = [m for m in available_models if 'gemini-2' in m][0]
    
    return {
        'aliases': aliases,
        'available_models': available_models,
        'recommended': recommended,
        'source': source
    }

def get_deepseek_models() -> dict:
    """
    Get a dictionary of available DeepSeek models for use with openai_api.
    
    This function returns information about DeepSeek models that can be used
    with the openai_api function in bilbystats. It fetches the current models
    from the DeepSeek API.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    use_api : bool, optional (default=True)
        Whether to fetch models from the DeepSeek API. If False, returns empty list.
    
    ---------------------------------------------------------------------------
    OUTPUT:
    models : dict
        A dictionary with the following structure:
        {
            'aliases': dict - Model aliases that openai_api accepts for DeepSeek
            'available_models': list - Full list of DeepSeek model identifiers
            'recommended': dict - Recommended models for different use cases
            'source': str - Whether models came from 'api' or 'unavailable'
        }
    
    ---------------------------------------------------------------------------
    AUTHORS: Your Name
    ---------------------------------------------------------------------------
    """
    import bilbystats as bs
    from openai import OpenAI
    
    # Model aliases for DeepSeek
    aliases = {
        "deepseek": "deepseek-chat",
    }
    
    available_models = []
    source = "unavailable"
    
    try:
            # Try to get models from the DeepSeek API
            key = bs.read_api_key("deepseek")
            client = OpenAI(
                api_key=key,
                base_url="https://api.deepseek.com"
            )
            
            # Get models from DeepSeek API
            models_response = client.models.list()
            available_models = [model.id for model in models_response.data]
            source = "api"
            print(f"✅ Fetched {len(available_models)} models from DeepSeek API")
                
    except Exception as e:
            print(f"⚠️  Could not fetch models from DeepSeek API: {e}")
            available_models = []
            source = "error"
    
    # Try to identify recommended models from the available list
    recommended = {}
    
    if available_models:
        # Look for specific patterns in available models
        chat_models = [m for m in available_models if 'chat' in m.lower()]
        coder_models = [m for m in available_models if 'coder' in m.lower()]
        reasoner_models = [m for m in available_models if 'reasoner' in m.lower() or 'r1' in m.lower()]
        
        if chat_models:
            recommended['balanced'] = sorted(chat_models)[-1]
            recommended['current_default'] = 'deepseek-chat'
        
        if coder_models:
            recommended['coding'] = sorted(coder_models)[-1]
        
        if reasoner_models:
            recommended['reasoning'] = sorted(reasoner_models)[-1]
            recommended['most_capable'] = sorted(reasoner_models)[-1]
        
        # Cost effective (DeepSeek is generally cost-effective)
        if available_models:
            recommended['cost_effective'] = available_models[0]
    
    return {
        'aliases': aliases,
        'available_models': available_models,
        'recommended': recommended,
        'source': source
    }

def get_models(provider: str) -> None:
    """
    Get and display model information in a pretty format for the specified provider.
    
    This function calls the appropriate get_*_models() function based on the provider
    and displays the results in a formatted, readable way.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
    provider : str
        The model provider. Must be one of: "deepseek", "openai", "gemini", "claude"
    
    ---------------------------------------------------------------------------
    OUTPUT:
    None
        Prints formatted model information to console
    
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    import bilbystats as bs
    
    # Map providers to their respective functions
    provider_functions = {
        "deepseek": bs.get_deepseek_models,
        "openai": bs.get_openai_models, 
        "gemini": bs.get_gemini_models,
        "claude": bs.get_claude_models
    }
    
    # Validate provider
    provider = provider.lower()
    if provider not in provider_functions:
        print(f"❌ Error: '{provider}' is not a valid provider.")
        print(f"   Valid providers: {', '.join(provider_functions.keys())}")
        return
    
    # Get model information
    try:
        models_info = provider_functions[provider]()
    except Exception as e:
        print(f"❌ Error getting models for {provider}: {e}")
        return
    
    # Pretty print the results
    print(f"\n{'='*60}")
    print(f"🤖 {provider.upper()} MODELS")
    print(f"{'='*60}")
    
    # Source information
    source_emoji = "🌐" if models_info['source'] == "api" else "⚠️"
    print(f"{source_emoji} Data source: {models_info['source']}")
    
    # Aliases section
    if provider != "openai":
        if models_info['aliases']:
            print(f"\n📝 ALIASES:")
            for alias, full_name in models_info['aliases'].items():
                print(f"   • {alias} → {full_name}")
    
    # Recommended models section
    #if models_info['recommended']:
    #    print(f"\n⭐ RECOMMENDED MODELS:")
    #    for use_case, model in models_info['recommended'].items():
    #        use_case_formatted = use_case.replace('_', ' ').title()
    #        print(f"   • {use_case_formatted}: {model}")
    
    # Available models section
    if models_info['available_models']:
        print(f"\n📋 ALL AVAILABLE MODELS ({len(models_info['available_models'])}):")
        #print(models_info['available_models'])
        for i, model in enumerate(sorted(models_info['available_models']), 1):
            print(f"   {i:2d}. {model}")
    else:
        print(f"\n📋 AVAILABLE MODELS: None found")
    
    print(f"{'='*80}\n")
    
def available_models(provider):
    available_models = bs.read_data('api_model_dict.json')
    print(f"\n📋 ALL AVAILABLE MODELS ({len(available_models[provider])}):")
    #print(models_info['available_models'])
    for i, model in enumerate(sorted(available_models[provider]), 1):
        print(f"   {i:2d}. {model}")