"""
Provider-agnostic LLM service using LiteLLM.
Supports Google Gemini, Anthropic Claude, OpenAI GPT, and OpenAI-compatible endpoints.
"""

import hashlib
import logging
from dataclasses import dataclass
from typing import Optional
import httpx
import litellm
from app.config.config import (
    AI_FIX_MOCK_MODE,
    AI_FALLBACK_MODEL_GOOGLE,
    AI_FALLBACK_MODEL_ANTHROPIC,
    AI_FALLBACK_MODEL_OPENAI,
)

logger = logging.getLogger(__name__)

# Disable litellm telemetry — prevents scan data/prompts from being sent to litellm servers
litellm.telemetry = False

# Fallback models used only when the provider's live API cannot be reached.
# Values come from environment variables — see app/config/config.py.
FALLBACK_MODELS: dict[str, str] = {
    "google":    AI_FALLBACK_MODEL_GOOGLE,
    "anthropic": AI_FALLBACK_MODEL_ANTHROPIC,
    "openai":    AI_FALLBACK_MODEL_OPENAI,
}

# In-memory cache: (provider, api_key_hash) -> model_id
# Populated on first generate call, cleared on server restart.
_model_cache: dict[str, str] = {}


@dataclass
class LLMConfig:
    """Configuration for LLM provider."""
    provider: str  # 'google', 'anthropic', 'openai', 'custom'
    api_key: Optional[str] = None
    base_url: Optional[str] = None  # For custom/on-premise endpoints
    model: Optional[str] = None     # Required for custom; auto-resolved for others


def resolve_model(config: LLMConfig) -> str:
    """
    Return the best available model for the provider by querying their live API.
    Result is cached in-memory for the lifetime of the server process.
    Falls back to FALLBACK_MODELS if the API call fails.

    For custom provider, the model must be explicitly set by the caller.
    """
    # Custom provider: model is always explicitly specified by the user
    if config.provider == "custom":
        if config.model:
            return config.model
        logger.warning("Custom provider requires a model name; no model specified")
        raise ValueError("Custom provider requires a model name (e.g. ollama/llama3, azure/gpt-4o)")

    # Standard providers: explicit model overrides auto-discovery
    if config.model:
        return config.model

    if not config.api_key:
        return FALLBACK_MODELS.get(config.provider, "gemini/gemini-flash-latest")

    key_hash = hashlib.sha256(config.api_key.encode()).hexdigest()[:16]
    cache_key = f"{config.provider}:{key_hash}"

    if cache_key in _model_cache:
        return _model_cache[cache_key]

    try:
        model = _fetch_best_model(config)
        if model:
            logger.info(f"Resolved model for {config.provider}: {model}")
            _model_cache[cache_key] = model
            return model
    except Exception as e:
        logger.warning(f"Model discovery failed for {config.provider}, using fallback: {e}")

    fallback = FALLBACK_MODELS.get(config.provider, "gemini/gemini-flash-latest")
    logger.info(f"Using fallback model for {config.provider}: {fallback}")
    return fallback


def _fetch_best_model(config: LLMConfig) -> Optional[str]:
    """Query the provider's models API and return the best model ID."""

    if config.provider == "google":
        resp = httpx.get(
            "https://generativelanguage.googleapis.com/v1beta/models",
            params={"key": config.api_key},
            timeout=10,
        )
        resp.raise_for_status()
        for m in resp.json().get("models", []):
            if "generateContent" in m.get("supportedGenerationMethods", []):
                raw = m.get("name", "")  # "models/gemini-2.0-flash"
                return "gemini/" + raw.removeprefix("models/")

    elif config.provider == "anthropic":
        resp = httpx.get(
            "https://api.anthropic.com/v1/models",
            headers={"x-api-key": config.api_key, "anthropic-version": "2023-06-01"},
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json().get("data", [])
        if data:
            return f"anthropic/{data[0]['id']}"

    elif config.provider == "openai":
        resp = httpx.get(
            "https://api.openai.com/v1/models",
            headers={"Authorization": f"Bearer {config.api_key}"},
            timeout=10,
        )
        resp.raise_for_status()
        prefixes = ("gpt-4", "o1", "o3", "o4")
        models = [
            m for m in resp.json().get("data", [])
            if m["id"].startswith(prefixes)
        ]
        if models:
            # Sort by creation timestamp descending to get the newest model
            models.sort(key=lambda m: m.get("created", 0), reverse=True)
            return f"openai/{models[0]['id']}"

    return None


def generate_completion(prompt: str, config: LLMConfig) -> str:
    """
    Generate text completion using LiteLLM.

    Args:
        prompt: The prompt to send to the LLM
        config: LLM configuration with provider, api_key, base_url

    Returns:
        Generated text response

    Raises:
        Exception: If LLM call fails (invalid API key, model not found, etc.)
    """
    # Mock mode for testing
    if AI_FIX_MOCK_MODE:
        logger.info("AI_FIX_MOCK_MODE enabled; returning mock response")
        return _mock_completion(prompt)

    model = resolve_model(config)

    kwargs = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
    }

    if config.api_key:
        kwargs["api_key"] = config.api_key

    if config.base_url:
        kwargs["api_base"] = config.base_url

    logger.info(f"Calling LLM: provider={config.provider}, model={model}")

    try:
        response = litellm.completion(**kwargs)
        return response.choices[0].message.content
    except Exception as e:
        logger.error(f"LLM completion failed: {str(e)}")
        raise


def _mock_completion(prompt: str) -> str:
    """Return mock response for testing."""
    return """--- a/test.py
+++ b/test.py
@@ -1,3 +1,5 @@
+# Security Fix: Use SHA-256 instead of MD5
+# MD5 is cryptographically broken and vulnerable to collisions
 import hashlib
-password = hashlib.md5(b"password").hexdigest()
+password = hashlib.sha256(b"password").hexdigest()
"""
