import json
import logging
from typing import Dict
from core.llm import chat

logger = logging.getLogger(__name__)

def route_query(query: str, active_modes: list = None) -> Dict:
    """
    Agentic Orchestrator: Uses a fast LLM to decide the exact route for the user's query.
    Returns a dictionary with routing decision.
    """
    if active_modes is None:
        active_modes = ["chat", "rag", "graph", "vision", "prediction", "hybrid", "chart", "image", "ml_pipeline"]
        
    prompt = f"""You are the Master Orchestrator for an AI Data Analysis Platform.
Your job is to read the user's query and route it to the CORRECT specialized agent.

AVAILABLE ROUTES:
- "chat": For simple greetings, thanking, small talk ("hi", "who are you", "thanks").
- "rag": For specific questions about the data documents ("what is the revenue in Q1?", "who are the top customers?").
- "chart": If the user explicitly asks for a chart, graph, visualization, or plot.
- "image": If the user explicitly asks for an image, img, png, or picture of a dashboard.
- "vision": If the user says something like "look at this image" or "what is in this picture".
- "prediction": If the user asks for a forecast, prediction, or future trend.
- "ml_pipeline": If the user asks to "train a model", "build a model", "predict churn", "run machine learning", etc.
- "etl": If the user asks to "clean data", "generate etl", "fix missing values", or "write python script to clean".

USER QUERY: "{query}"

Determine the best route.
Output ONLY a valid JSON object in this exact format:
{{
    "route": "one_of_the_routes_above",
    "confidence": 95,
    "reason": "Brief explanation why"
}}
"""
    try:
        # Use a fast, reliable model for orchestration (e.g., Llama 3 8B)
        # Using temperature=0.1 for deterministic routing
        response = chat(prompt, temperature=0.1, max_tokens=150)
        
        # Parse JSON
        response = response.strip()
        if '```json' in response:
            response = response.split('```json')[1].split('```')[0]
        elif '```' in response:
            response = response.split('```')[1].split('```')[0]
            
        start = response.find('{')
        end = response.rfind('}') + 1
        if start >= 0 and end > start:
            response = response[start:end]
            
        decision = json.loads(response)
        
        # Validate route
        valid_routes = ["chat", "rag", "graph", "vision", "prediction", "hybrid", "chart", "image", "etl", "ml_pipeline"]
        if decision.get("route") not in valid_routes:
            decision["route"] = "rag"  # fallback
            
        logger.info(f"Orchestrator routed '{query}' -> {decision['route']} (Confidence: {decision.get('confidence')}%)")
        return decision
        
    except Exception as e:
        logger.error(f"Orchestrator failed: {e}")
        # Fallback logic if LLM routing fails
        query_lower = query.lower()
        if any(kw in query_lower for kw in ['img', 'image', 'picture', 'dashboard']):
            return {"route": "image", "confidence": 50, "reason": "Fallback: Keyword match"}
        if any(kw in query_lower for kw in ['chart', 'graph', 'plot', 'visualize', 'trend']):
            return {"route": "chart", "confidence": 50, "reason": "Fallback: Keyword match"}
        if any(kw in query_lower for kw in ['predict', 'forecast', 'future']):
            return {"route": "prediction", "confidence": 50, "reason": "Fallback: Keyword match"}
        if any(kw in query_lower for kw in ['clean', 'fix', 'etl', 'python script to clean']):
            return {"route": "etl", "confidence": 50, "reason": "Fallback: Keyword match"}
        if any(kw in query_lower for kw in ['hi', 'hello', 'hey', 'thanks', 'bye']):
            return {"route": "chat", "confidence": 50, "reason": "Fallback: Keyword match"}
            
        return {"route": "rag", "confidence": 50, "reason": "Fallback: Default route"}