services/ai_router.py

import httpx
import json
import re
import logging
from config.settings import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()


def _extract_json(text: str) -> dict:
    """
    Bulletproof JSON extraction for Level 4-8 'Heavyweight' models.
    Supports markdown fences, meta-commentary, and nested structures.
    """
    text = text.strip()
    
    # 1. Trial: Pure JSON
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    
    # 2. Trial: Regex for ```json codes blocks
    fence_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", text)
    if fence_match:
        try:
            return json.loads(fence_match.group(1).strip())
        except json.JSONDecodeError:
            pass

    # 3. Trial: Brute-force first '{' to last '}'
    start = text.find('{')
    end = text.rfind('}')
    if start != -1 and end != -1 and end > start:
        json_str = text[start:end+1]
        try:
            return json.loads(json_str)
        except json.JSONDecodeError:
            pass

    raise ValueError(f"Could not extract valid JSON from AI response. Status: {text[:100]}...")

def _sanitize_result(data: dict) -> dict:
    """
    Ensures that 'issues' and 'suggestions' are always flat lists of strings.
    Handles 'Mistral Categorization' where the AI returns objects instead of arrays.
    """
    for key in ["issues", "suggestions"]:
        val = data.get(key)
        if isinstance(val, dict):
            # Flatten dictionary: {"category": "issue"} -> ["category: issue"]
            new_list = []
            for k, v in val.items():
                if isinstance(v, list):
                    new_list.extend([f"{k}: {item}" for item in v])
                else:
                    new_list.append(f"{k}: {v}")
            data[key] = new_list
        elif val and not isinstance(val, list):
            data[key] = [str(val)]
        elif not val:
            data[key] = []
            
    return data


# Smart Model Dictionary: Maps logical names to provider-specific model IDs
# Ordered from Level 1 (Entry) to Level 8 (Flagship)
MODEL_MAP = {
    "gemma-4-31b": {
        "openrouter": "google/gemma-4-31b-it:free"
    },
    "llama-3.1": {
        "groq": "llama-3.1-8b-instant",
        "openrouter": "meta-llama/llama-3.1-8b-instruct:free"
    },
    "qwen-2.5": {
        "openrouter": "qwen/qwen-2.5-7b-instruct:free",
        "huggingface": "Qwen/Qwen2.5-7B-Instruct"
    },
    "nemotron-120b": {
        "openrouter": "nvidia/nemotron-3-super-120b-a12b:free"
    },
    "minimax-2.5": {
        "openrouter": "minimax/minimax-m2.5:free"
    },
    "mistral-large": {
        "openrouter": "mistralai/mistral-large-2407"
    },
    "groq-70b": {
        "groq": "llama-3.3-70b-versatile",
        "openrouter": "meta-llama/llama-3.3-70b-instruct:free"
    },
    "gemini-flash": {
        "gemini": "gemini-2.0-flash",
        "openrouter": "google/gemini-2.0-flash-001"
    }
}


async def analyze_with_groq(prompt: str, model_id: str = "llama-3.1-8b-instant") -> dict:
    """Backup AI — Groq with specified model."""
    if not settings.groq_api_key:
        raise Exception("Groq API key not configured")
        
    # Dynamic timeout for large codebases
    timeout = 60.0 if len(prompt) > 5000 else 30.0
    
    async with httpx.AsyncClient(timeout=timeout) as client:
        response = await client.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {settings.groq_api_key}", "Content-Type": "application/json"},
            json={
                "model": model_id, 
                "messages": [{"role": "user", "content": prompt}], 
                "temperature": 0.1,
                "max_tokens": 2000 # Increased for large code analysis
            }
        )
        if response.status_code != 200:
            raise Exception(f"Groq error {response.status_code}: {response.text[:200]}")
        result = response.json()
        text = result["choices"][0]["message"]["content"].strip()
        return _sanitize_result(_extract_json(text))


async def analyze_with_gemini(prompt: str, model_id: str = "gemini-2.0-flash") -> dict:
    """Primary AI — Gemini with specified model."""
    if not settings.gemini_api_key:
        raise Exception("Gemini API key not configured")
        
    # Support large context in Gemini
    timeout = 60.0 if len(prompt) > 5000 else 30.0
    
    async with httpx.AsyncClient(timeout=timeout) as client:
        response = await client.post(
            f"https://generativelanguage.googleapis.com/v1beta/models/{model_id}:generateContent?key={settings.gemini_api_key}",
            json={
                "contents": [{"parts": [{"text": prompt}]}],
                "generationConfig": {"maxOutputTokens": 2000} # Increased
            }
        )
        if response.status_code != 200:
            raise Exception(f"Gemini error {response.status_code}: {response.text[:200]}")
        result = response.json()
        candidates = result.get("candidates")
        if not candidates:
            raise Exception("Gemini returned no candidates")
        text = candidates[0]["content"]["parts"][0]["text"].strip()
        return _sanitize_result(_extract_json(text))


async def analyze_with_openrouter(prompt: str, model_id: str = None) -> dict:
    """Last Resort AI — OpenRouter with customizable model."""
    if not settings.openrouter_api_key:
        raise Exception("OpenRouter API key not configured")
    
    # Use settings default if no specific model requested
    target_model = model_id or settings.openrouter_model
        
    # Support large context in OpenRouter
    timeout = 60.0 if len(prompt) > 5000 else 30.0
    
    async with httpx.AsyncClient(timeout=timeout) as client:
        response = await client.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers={"Authorization": f"Bearer {settings.openrouter_api_key}", "Content-Type": "application/json"},
            json={
                "model": target_model, 
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 2000 # Increased
            }
        )
        if response.status_code != 200:
            raise Exception(f"OpenRouter error {response.status_code}: {response.text[:200]}")
        result = response.json()
        if "choices" not in result:
            raise Exception(f"Unexpected OpenRouter response: {result}")
        text = result["choices"][0]["message"]["content"].strip()
        return _sanitize_result(_extract_json(text))


async def analyze_with_huggingface(prompt: str, model_id: str = None) -> dict:
    """Extra Backup — Hugging Face with specific model OR rotation."""
    if not settings.huggingface_api_token:
        raise Exception("Hugging Face API token not configured")
        
    # If a specific model is requested, try ONLY that one
    # Otherwise, use the rotation logic
    models_to_try = [model_id] if model_id else settings.huggingface_free_models
    
    last_error = None
    for target_model in models_to_try:
        try:
            timeout = 60.0 if len(prompt) > 5000 else 30.0
            async with httpx.AsyncClient(timeout=timeout) as client:
                response = await client.post(
                    f"https://api-inference.huggingface.co/models/{target_model}",
                    headers={"Authorization": f"Bearer {settings.huggingface_api_token}"},
                    json={
                        "inputs": prompt,
                        "parameters": {"max_new_tokens": 2000, "return_full_text": False}
                    }
                )
                if response.status_code != 200:
                    raise Exception(f"HF Model {target_model} failed ({response.status_code})")
                
                result = response.json()
                text = result[0]["generated_text"].strip() if isinstance(result, list) else result.get("generated_text", "").strip()
                return _sanitize_result(_extract_json(text))
        except Exception as e:
            logger.warning(f"Hugging Face model {target_model} failed: {e}")
            last_error = e
            if model_id: # Don't rotate if specific model was forced
                break
            continue

    raise Exception(f"Hugging Face attempt failed. Last error: {last_error}")


async def route_analysis(prompt: str, model_choice: str = "auto") -> dict:
    """
    Smart AI router with Zero-Failure 'Smart Recovery' logic.
    
    Tiered Chain: Gemma 4 -> Llama 3.1 -> Qwen 2.5 -> Nemotron 120B -> 
                 MiniMax 2.5 -> Mistral Large -> Groq 70B -> Gemini Flash
    """
    
    # 1. Prepare Provider List
    providers = []
    
    # Handle Specific Model Choice (With Recovery)
    if model_choice != "auto" and model_choice in MODEL_MAP:
        config = MODEL_MAP[model_choice]
        if "groq" in config:
            providers.append({"func": lambda p: analyze_with_groq(p, config["groq"]), "name": model_choice})
        if "gemini" in config:
            providers.append({"func": lambda p: analyze_with_gemini(p, config["gemini"]), "name": model_choice})
        if "openrouter" in config:
            providers.append({"func": lambda p: analyze_with_openrouter(p, config["openrouter"]), "name": model_choice})
        if "huggingface" in config:
            providers.append({"func": lambda p: analyze_with_huggingface(p, config["huggingface"]), "name": model_choice})
        
        # SMART RECOVERY: If the specific choice fails, pivot to the full AUTO chain
        providers.append({"func": lambda p: route_analysis(p, "auto"), "is_meta": True})
        
    else:
        # Full Power Progression Chain (8 Levels)
        order = ["gemma-4-31b", "llama-3.1", "qwen-2.5", "nemotron-120b", "minimax-2.5", "mistral-large", "groq-70b", "gemini-flash"]
        for m_id in order:
            m_cfg = MODEL_MAP[m_id]
            # Primary provider logic for auto-chain
            if m_id == "llama-3.1": # Prioritize Groq's 8B for speed
                providers.append({"func": lambda p: analyze_with_groq(p, "llama-3.1-8b-instant"), "name": m_id})
            elif m_id == "groq-70b": # Prioritize Groq's 70B
                providers.append({"func": lambda p: analyze_with_groq(p, "llama-3.3-70b-versatile"), "name": m_id})
            elif m_id == "gemini-flash":
                providers.append({"func": lambda p: analyze_with_gemini(p), "name": m_id})
            elif "openrouter" in m_cfg:
                providers.append({"func": lambda p: analyze_with_openrouter(p, m_cfg["openrouter"]), "name": m_id})
        
        # Ultimate Last Resort
        providers.append({"func": lambda p: analyze_with_openrouter(p, "openrouter/free"), "name": "openrouter-free-fallback"})

    # 2. Execute with Failover Logic
    last_error = None
    for i, p_item in enumerate(providers):
        try:
            p_name = p_item.get("name", "failover-chain")
            logger.info(f"Attempting {p_name} ({i+1}/{len(providers)})")
            
            result = await p_item["func"](prompt)
            
            # If it's a metadata-wrapped result from a nested call, return it
            if isinstance(result, dict) and "_actual_model" in result:
                return result
                
            # Attach actual model name to the result
            if isinstance(result, dict):
                result["_actual_model"] = p_name
            return result
            
        except Exception as e:
            logger.error(f"Provider {i+1} failed: {str(e)}")
            last_error = e
            continue

    raise Exception(f"Absolute failure in AI Hub. Last error: {last_error}")