import json
import os
import re
from litellm import completion
from src.config import Config
from src.crews.definitions import CrewDefinitions

class SmartRouter:
    """
    Decides which Crew should handle a user request using a Fast LLM.
    """
    
    @staticmethod
    def route(user_input: str) -> str:
        """
        Returns the name of the Crew that should handle the input.
        """
        
        # Get available crews
        crews = CrewDefinitions.get_available_crews()
        crews_str = ", ".join([f"'{c}'" for c in crews])
        
        prompt = f"""
        You are an intelligent dispatcher for an IT Enterprise system.
        
        AVAILABLE CREWS:
        1. 'Infra Engineering (Zabbix)': For monitoring, server health, Zabbix templates, UUIDs, YAML validation.
        2. 'Security Audit': For hacking, vulnerability checks, penetration testing.
        3. 'HR & Evolution': For creating new agents, learning new rules, system updates.
        4. 'Sales Growth': For sales pitches, CRM, deals, negotiation.
        5. 'Business Strategy': For ROI, legal, compliance, high-level decisions.
        
        USER INPUT: "{user_input}"
        
        TASK:
        Classify the input into exactly one of the available crews.
        Return ONLY valid JSON like this: {{"crew": "Infra Engineering (Zabbix)"}}
        If unsure, default to "Business Strategy" (Generalist).
        Do NOT include any text before or after the JSON.
        """
        
        config = Config.get_llm_config(mode="fast")
        provider = os.getenv("LLM_PROVIDER", "openai").lower()
        
        try:
            # Build completion kwargs - only add response_format for OpenAI
            completion_kwargs = {
                "model": config['model'],
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.0,
            }
            
            # Add optional params based on provider
            if config.get('api_key'):
                completion_kwargs["api_key"] = config['api_key']
            if config.get('base_url'):
                completion_kwargs["base_url"] = config['base_url']
            
            # response_format only works with OpenAI-compatible APIs
            if provider == "openai":
                completion_kwargs["response_format"] = {"type": "json_object"}
            
            response = completion(**completion_kwargs)
            
            content = response.choices[0].message.content.strip()
            
            # Robust JSON parsing (handling markdown blocks and extra text)
            # Try to extract JSON from the response
            json_match = re.search(r'\{[^}]+\}', content)
            if json_match:
                content = json_match.group(0)
            elif "```json" in content:
                content = content.split("```json")[1].split("```")[0].strip()
            elif "```" in content:
                content = content.split("```")[1].split("```")[0].strip()
                
            data = json.loads(content)
            selected = data.get("crew", "Business Strategy")
            
            # Validation
            if selected not in crews:
                return "Business Strategy"
                
            return selected

        except Exception as e:
            print(f"Routing Error: {e}")
            return "Business Strategy" # Fallback