#!/usr/bin/env python3
"""
Check OpenRouter model prices and print free models
"""
import requests
import json
import os
import logging
import shutil

from openrouter_client import get_api_key

# ANSI escape codes for styling
GRAY_ITALIC = "\033[3;90m"  # Italic + bright black (light gray)
RESET = "\033[0m"

logger = logging.getLogger(__name__)

OPENROUTER_MODELS_URL = "https://openrouter.ai/api/v1/models/user"


def fetch_models():
    """
    Fetch available models from OpenRouter API

    Returns:
        tuple: (success: bool, models: list or None, error: str or None)
    """
    api_key = get_api_key()

    if not api_key:
        logger.error("OPENROUTER_API_KEY environment variable not set and ~/.openai_openrouter.key not found")
        return False, None, "Missing API key"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    try:
        logger.info("Fetching models from OpenRouter API...")
        response = requests.get(OPENROUTER_MODELS_URL, headers=headers, timeout=30)
        response.raise_for_status()

        data = response.json()
        models = data.get("data", [])

        logger.info(f"Successfully fetched {len(models)} models")
        return True, models, None

    except requests.exceptions.Timeout:
        logger.error("OpenRouter API request timed out")
        return False, None, "Timeout"

    except requests.exceptions.RequestException as e:
        logger.error(f"OpenRouter API request failed: {e}")
        return False, None, str(e)

    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse API response: {e}")
        return False, None, f"JSON decode error: {e}"

    except Exception as e:
        logger.error(f"Unexpected error fetching models: {e}", exc_info=True)
        return False, None, str(e)


def filter_free_models(models):
    """
    Filter models that are free (pricing is 0)

    Args:
        models: List of model dictionaries

    Returns:
        list: Free model dictionaries
    """
    free_models = []

    for model in models:
        pricing = model.get("pricing", {})
        prompt_price = pricing.get("prompt", "0")
        completion_price = pricing.get("completion", "0")

        if prompt_price == "0" and completion_price == "0":
            free_models.append(model)

    return free_models


def extract_model_size(model_id):
    """
    Extract model size from model ID string

    Args:
        model_id: Model ID string (e.g., "openai/gpt-4-30b:free")

    Returns:
        int: Model size in billions, or 0 if not found
    """
    import re

    pattern = r'-e?(\d+)b[:\-]'
    match = re.search(pattern, model_id.lower())

    if match:
        return int(match.group(1))

    return 0


def print_free_models(free_models):
    """
    Print free models in a readable format, sorted by model size (descending)

    Args:
        free_models: List of free model dictionaries
    """
    if not free_models:
        print("No free models found")
        return

    sorted_models = sorted(free_models, key=lambda m: extract_model_size(m.get('id', '')), reverse=True)

    # Find the longest ID line to align Size/Context column
    max_id_len = 0
    for model in sorted_models:
        model_id = model.get('id', 'Unknown')
        id_line = f"{model_id}"
        max_id_len = max(max_id_len, len(id_line))

    # Column where Size/Context starts (add some padding)
    size_col = max_id_len + 8  # "N. ID: " is ~7 chars, plus 1 space padding

    print(f"\nFound {len(sorted_models)} free models (sorted by model size):\n")
    print("-" * 100)

    for i, model in enumerate(sorted_models, 1):
        model_id = model.get('id', 'Unknown')
        size = extract_model_size(model_id)
        size_str = f"{size}B" if size > 0 else "Unknown"
        context = model.get('context_length', 'Unknown')

        # Build the ID line
        id_text = f"{i}. ID: {model_id}"
        # Calculate padding needed to reach size_col
        padding = max(1, size_col - len(id_text))
        # Print ID and Size/Context on same line
        print(f"{id_text}{' ' * padding}Size/Context: {size_str}/{context}")
        # Print description on next line (remove newlines, fill terminal width)
        terminal_width = shutil.get_terminal_size().columns
        # Account for "   " prefix (3 spaces)
        desc_max_len = terminal_width - 3
        raw_description = model.get('description', 'No description').replace('\n', ' ')
        description = raw_description[:desc_max_len]
        print(f"   {GRAY_ITALIC}{description}{RESET}")

    print("-" * 100)


def filter_paid_models(models):
    """
    Filter models that are paid (pricing is not 0)

    Args:
        models: List of model dictionaries

    Returns:
        list: Paid model dictionaries
    """
    paid_models = []

    for model in models:
        pricing = model.get("pricing", {})
        prompt_price = pricing.get("prompt", "0")
        completion_price = pricing.get("completion", "0")

        # Consider a model paid if either prompt or completion price is non-zero
        if prompt_price != "0" or completion_price != "0":
            paid_models.append(model)

    return paid_models


def print_paid_models(paid_models):
    """
    Print paid models in a readable format, sorted by model size (descending)

    Args:
        paid_models: List of paid model dictionaries
    """
    if not paid_models:
        print("No paid models found")
        return

    sorted_models = sorted(paid_models, key=lambda m: m.get('created', 0), reverse=True)

    # Find the longest ID line to align Size/Context column
    max_id_len = 0
    for model in sorted_models:
        model_id = model.get('id', 'Unknown')
        id_line = f"{model_id}"
        max_id_len = max(max_id_len, len(id_line))

    # Column where Size/Context starts (add some padding)
    size_col = max_id_len + 8  # "N. ID: " is ~7 chars, plus 1 space padding

    print(f"\nFound {len(sorted_models)} paid models (sorted by newest):\n")
    print("-" * 100)

    for i, model in enumerate(sorted_models[:50], 1):  # Limit to first 50 to avoid spam
        model_id = model.get('id', 'Unknown')
        size = extract_model_size(model_id)
        size_str = f"{size}B" if size > 0 else "Unknown"
        context_raw = model.get('context_length', 'Unknown')
        # Convert context to k tokens
        if isinstance(context_raw, (int, float)) and context_raw != 'Unknown':
            context_k = context_raw / 1000
            if context_k >= 1000:
                context = f"{context_k/1000:.1f}M"
            else:
                context = f"{context_k:.0f}k"
        else:
            context = str(context_raw)
        pricing = model.get('pricing', {})
        prompt_price_raw = pricing.get('prompt', '0')
        completion_price_raw = pricing.get('completion', '0')
        # Convert to per 1M tokens
        try:
            prompt_price_mil = float(prompt_price_raw) * 1_000_000
            completion_price_mil = float(completion_price_raw) * 1_000_000
            # Format: if < 1, show 4 decimals; if < 10, show 2 decimals; else show 0 decimals
            if prompt_price_mil < 1:
                prompt_str = f"${prompt_price_mil:.4f}"
            elif prompt_price_mil < 10:
                prompt_str = f"${prompt_price_mil:.2f}"
            else:
                prompt_str = f"${prompt_price_mil:.0f}"
            
            if completion_price_mil < 1:
                completion_str = f"${completion_price_mil:.4f}"
            elif completion_price_mil < 10:
                completion_str = f"${completion_price_mil:.2f}"
            else:
                completion_str = f"${completion_price_mil:.0f}"
            
            price_str = f"{prompt_str}/{completion_str}"
        except (ValueError, TypeError):
            price_str = f"${prompt_price_raw}/${completion_price_raw}"

        # Build the ID line
        id_text = f"{i}. ID: {model_id}"
        # Calculate padding needed to reach size_col
        padding = max(1, size_col - len(id_text))
        # Print ID and Size/Context/Price on same line
        print(f"{id_text}{' ' * padding}Size/Context(k)/Price: {size_str}/{context}/{price_str}")
        # Print description on next line (remove newlines, fill terminal width)
        terminal_width = shutil.get_terminal_size().columns
        # Account for "   " prefix (3 spaces)
        desc_max_len = terminal_width - 3
        raw_description = model.get('description', 'No description').replace('\n', ' ')
        description = raw_description[:desc_max_len]
        print(f"   {GRAY_ITALIC}{description}{RESET}")

    if len(sorted_models) > 50:
        print(f"\n   ... and {len(sorted_models) - 50} more paid models")

    print("-" * 100)


def main():
    """
    Main function to fetch and display free models
    """
    logging.basicConfig(level=logging.INFO)

    success, models, error = fetch_models()

    if not success:
        print(f"Error: {error}")
        return

    free_models = filter_free_models(models)
    print_free_models(free_models)

    print(f"\nTotal models: {len(models)}")
    print(f"Free models: {len(free_models)}")


if __name__ == "__main__":
    main()
