"""
AI Provider Configuration
Central configuration for AI backends (Groq and g4f)
Controls which backend is used globally and per-model
"""

import os
from dotenv import load_dotenv

load_dotenv()

# ---------------------------------------------------------------------------
# Global AI Backend Toggle
# Values: 'groq' (default, safe & fast), 'g4f' (free/scraping), 'auto' (groq with g4f fallback)
# ---------------------------------------------------------------------------
AI_BACKEND = os.getenv('AI_BACKEND', 'groq')

# ---------------------------------------------------------------------------
# Groq Configuration
# ---------------------------------------------------------------------------
GROQ_API_KEY = os.getenv('GROQ_API_KEY', '')
GROQ_BASE_URL = 'https://api.groq.com/openai/v1'

# Default models
GROQ_DEFAULT_MODEL = os.getenv('GROQ_DEFAULT_MODEL', 'llama-3.3-70b-versatile')
GROQ_VISION_MODEL = os.getenv('GROQ_VISION_MODEL', 'meta-llama/llama-4-scout-17b-16e-instruct')
GROQ_FAST_MODEL = os.getenv('GROQ_FAST_MODEL', 'llama-3.1-8b-instant')

# ---------------------------------------------------------------------------
# Groq Model Registry
# Maps friendly/generic names → exact Groq model IDs
# ---------------------------------------------------------------------------
GROQ_MODEL_MAPPING = {
    # Production models
    'llama-3.3-70b': 'llama-3.3-70b-versatile',
    'llama-3.3-70b-versatile': 'llama-3.3-70b-versatile',
    'llama-3.1-8b': 'llama-3.1-8b-instant',
    'llama-3.1-8b-instant': 'llama-3.1-8b-instant',
    'gpt-oss-20b': 'openai/gpt-oss-20b',
    'openai/gpt-oss-20b': 'openai/gpt-oss-20b',
    'gpt-oss-120b': 'openai/gpt-oss-120b',
    'openai/gpt-oss-120b': 'openai/gpt-oss-120b',

    # Vision model
    'llama-4-scout': 'meta-llama/llama-4-scout-17b-16e-instruct',
    'meta-llama/llama-4-scout-17b-16e-instruct': 'meta-llama/llama-4-scout-17b-16e-instruct',

    # Preview models
    'qwen3-32b': 'qwen/qwen3-32b',
    'qwen/qwen3-32b': 'qwen/qwen3-32b',
    'kimi-k2': 'moonshotai/kimi-k2-instruct-0905',
    'moonshotai/kimi-k2-instruct-0905': 'moonshotai/kimi-k2-instruct-0905',

    # Compound systems
    'compound': 'groq/compound',
    'groq/compound': 'groq/compound',
    'compound-mini': 'groq/compound-mini',
    'groq/compound-mini': 'groq/compound-mini',

    # Legacy/generic name aliases → best Groq equivalent
    'gpt-4o': 'llama-3.3-70b-versatile',
    'gpt-4o-mini': 'llama-3.1-8b-instant',
    'gpt-4': 'llama-3.3-70b-versatile',
    'gpt-3.5-turbo': 'llama-3.1-8b-instant',
    'gpt-5-mini': 'llama-3.1-8b-instant',
    'gemini-search': 'llama-3.3-70b-versatile',
    'geminisearch': 'llama-3.3-70b-versatile',
    'gemini': 'llama-3.3-70b-versatile',
    'gemini-pro': 'llama-3.3-70b-versatile',
    'gemini-fast': 'llama-3.1-8b-instant',
    'claude': 'llama-3.3-70b-versatile',
    'claude-3.5-sonnet': 'llama-3.3-70b-versatile',
    'claude-3-opus': 'openai/gpt-oss-120b',
    'claude-fast': 'llama-3.1-8b-instant',
    'deepseek-chat': 'llama-3.3-70b-versatile',
    'deepseek-coder': 'qwen/qwen3-32b',
    'deepseek-r1': 'openai/gpt-oss-120b',
    'openai': 'llama-3.3-70b-versatile',
    'openai-fast': 'llama-3.1-8b-instant',
    'openai-large': 'openai/gpt-oss-120b',
    'qwen-coder': 'qwen/qwen3-32b',
    'qwen-2.5-coder-32b': 'qwen/qwen3-32b',
    'mistral': 'llama-3.3-70b-versatile',
    'mistral-large': 'llama-3.3-70b-versatile',
    'grok': 'llama-3.3-70b-versatile',
    'grok-2': 'llama-3.3-70b-versatile',
    'command-r-plus-08-2024': 'llama-3.3-70b-versatile',  # streaming model alias
    'claude40opusthinking': 'llama-3.3-70b-versatile',     # notes_audio fake model
    'o4-mini': 'llama-3.1-8b-instant',
}

# ---------------------------------------------------------------------------
# Groq Model Capabilities
# ---------------------------------------------------------------------------
GROQ_MODEL_CAPABILITIES = {
    'llama-3.3-70b-versatile': {
        'supports_vision': False,
        'supports_tools': True,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 131072,
        'max_output_tokens': 32768,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 300000,
        'tokens_per_second': 280,
        'cost_input_per_1m': 0.59,
        'cost_output_per_1m': 0.79,
    },
    'llama-3.1-8b-instant': {
        'supports_vision': False,
        'supports_tools': False,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 131072,
        'max_output_tokens': 131072,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 250000,
        'tokens_per_second': 560,
        'cost_input_per_1m': 0.05,
        'cost_output_per_1m': 0.08,
    },
    'openai/gpt-oss-20b': {
        'supports_vision': False,
        'supports_tools': True,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 131072,
        'max_output_tokens': 65536,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 250000,
        'tokens_per_second': 1000,
        'cost_input_per_1m': 0.075,
        'cost_output_per_1m': 0.30,
    },
    'openai/gpt-oss-120b': {
        'supports_vision': False,
        'supports_tools': True,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 131072,
        'max_output_tokens': 65536,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 250000,
        'tokens_per_second': 500,
        'cost_input_per_1m': 0.15,
        'cost_output_per_1m': 0.60,
    },
    'meta-llama/llama-4-scout-17b-16e-instruct': {
        'supports_vision': True,
        'supports_tools': True,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 131072,
        'max_output_tokens': 8192,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 300000,
        'tokens_per_second': 750,
        'cost_input_per_1m': 0.11,
        'cost_output_per_1m': 0.34,
    },
    'qwen/qwen3-32b': {
        'supports_vision': False,
        'supports_tools': False,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 131072,
        'max_output_tokens': 40960,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 300000,
        'tokens_per_second': 400,
        'cost_input_per_1m': 0.29,
        'cost_output_per_1m': 0.59,
    },
    'moonshotai/kimi-k2-instruct-0905': {
        'supports_vision': False,
        'supports_tools': True,
        'supports_streaming': True,
        'supports_json_mode': True,
        'context_window': 262144,
        'max_output_tokens': 16384,
        'rate_limit_rpm': 1000,
        'rate_limit_tpm': 250000,
        'tokens_per_second': 200,
        'cost_input_per_1m': 1.00,
        'cost_output_per_1m': 3.00,
    },
    'groq/compound': {
        'supports_vision': False,
        'supports_tools': True,
        'supports_streaming': True,
        'supports_json_mode': False,
        'context_window': 131072,
        'max_output_tokens': 8192,
        'rate_limit_rpm': 200,
        'rate_limit_tpm': 200000,
        'tokens_per_second': 450,
        'cost_input_per_1m': 0,
        'cost_output_per_1m': 0,
    },
}


def resolve_groq_model(model_name: str) -> str:
    """Resolve any model name to a valid Groq model ID.
    Returns the mapped Groq model ID, or the default model if unknown."""
    if not model_name:
        return GROQ_DEFAULT_MODEL
    return GROQ_MODEL_MAPPING.get(model_name, GROQ_DEFAULT_MODEL)


def resolve_groq_vision_model() -> str:
    """Get the vision-capable Groq model."""
    return GROQ_VISION_MODEL


def get_groq_model_capabilities(model_id: str) -> dict:
    """Get capability info for a Groq model."""
    return GROQ_MODEL_CAPABILITIES.get(model_id, GROQ_MODEL_CAPABILITIES.get(GROQ_DEFAULT_MODEL, {}))


def is_groq_enabled() -> bool:
    """Check if Groq backend is enabled and configured."""
    return AI_BACKEND in ('groq', 'auto') and bool(GROQ_API_KEY)


def is_g4f_enabled() -> bool:
    """Check if g4f backend is enabled."""
    return AI_BACKEND in ('g4f', 'auto')
