#!/usr/bin/env python3
"""
ISL Router: Complexity-Based Task Routing
Routes tasks between local OLMoE and external tools based on ISL principles
"""

import re
from typing import Dict, List, Tuple
from enum import Enum

class TaskComplexity(Enum):
    """Task complexity levels for routing decisions"""
    TRIVIAL = 1      # Direct lookup, simple calc
    SIMPLE = 2       # Basic reasoning, OLMoE handles
    MODERATE = 3     # Needs OLMoE + maybe tools
    COMPLEX = 4      # Requires external specialist
    UNCERTAIN = 5    # Needs search/verification

class RoutingDecision(Enum):
    """Routing strategies"""
    LOCAL_ONLY = "local"           # Pure OLMoE
    LOCAL_WITH_TOOLS = "hybrid"    # OLMoE + tools
    DELEGATE = "external"          # External specialist
    SEARCH_FIRST = "search"        # Search then OLMoE

class ISLRouter:
    """
    ISL-Style Router for Hybrid MoE System
    
    Philosophy:
    - Measure complexity BEFORE calling expensive models
    - Delegate only when needed
    - Avoid over-thinking simple tasks
    - Log when OLMoE was enough
    """
    
    def __init__(self):
        # Complexity scoring patterns
        self.trivial_patterns = [
            r"what is \d+\s*[\+\-\*\/]\s*\d+",  # Simple math
            r"define \w+",                       # Definitions
            r"what time",                        # Time queries
        ]
        
        self.complex_patterns = [
            r"prove|theorem|lemma",              # Math proofs
            r"implement.*algorithm",             # Complex code
            r"analyze.*security",                # Security analysis
            r"compare.*benchmark",               # Deep comparisons
        ]
        
        self.search_patterns = [
            r"latest|recent|current|today",      # Time-sensitive
            r"who is|who was",                   # Entity queries
            r"news|breaking",                    # News
        ]
        
        # Tool availability
        self.available_tools = {
            'calculator': True,
            'search': True,
            'code_executor': True,
            'specialist_llm': False  # Expensive, use sparingly
        }
    
    def score_complexity(self, query: str) -> Tuple[TaskComplexity, float]:
        """
        Score task complexity using ISL heuristics
        
        Returns:
            (complexity_level, confidence_score)
        """
        query_lower = query.lower()
        
        # Check trivial patterns
        for pattern in self.trivial_patterns:
            if re.search(pattern, query_lower):
                return TaskComplexity.TRIVIAL, 0.9
        
        # Check search patterns
        for pattern in self.search_patterns:
            if re.search(pattern, query_lower):
                return TaskComplexity.UNCERTAIN, 0.8
        
        # Check complex patterns
        for pattern in self.complex_patterns:
            if re.search(pattern, query_lower):
                return TaskComplexity.COMPLEX, 0.85
        
        # Heuristics
        word_count = len(query.split())
        has_code = '```' in query or 'def ' in query or 'function' in query or 'write a' in query_lower
        has_math = any(op in query for op in ['+', '-', '*', '/', '=', '^']) or 'calculate' in query_lower
        
        # Scoring logic
        if word_count < 10 and not has_code:
            # Check if it's a math problem that needs calculator
            if has_math and 'calculate' in query_lower:
                return TaskComplexity.MODERATE, 0.8  # Trigger tools
            return TaskComplexity.SIMPLE, 0.7
        elif has_code or word_count > 50:
            return TaskComplexity.MODERATE, 0.6
        else:
            return TaskComplexity.SIMPLE, 0.65
    
    def route(self, query: str) -> Dict:
        """
        Make routing decision based on complexity
        
        Returns:
            {
                'strategy': RoutingDecision,
                'complexity': TaskComplexity,
                'confidence': float,
                'tools': List[str],
                'reasoning': str
            }
        """
        complexity, confidence = self.score_complexity(query)
        
        # Routing logic
        if complexity == TaskComplexity.TRIVIAL:
            return {
                'strategy': RoutingDecision.LOCAL_ONLY,
                'complexity': complexity,
                'confidence': confidence,
                'tools': [],
                'reasoning': 'Trivial task, OLMoE handles directly'
            }
        
        elif complexity == TaskComplexity.SIMPLE:
            return {
                'strategy': RoutingDecision.LOCAL_ONLY,
                'complexity': complexity,
                'confidence': confidence,
                'tools': [],
                'reasoning': 'Simple reasoning, OLMoE sufficient'
            }
        
        elif complexity == TaskComplexity.MODERATE:
            # Check if tools would help
            tools_needed = []
            if 'calculate' in query.lower() or any(op in query for op in ['+', '-', '*', '/']):
                tools_needed.append('calculator')
            if 'search' in query.lower() or 'find' in query.lower():
                tools_needed.append('search')
            
            return {
                'strategy': RoutingDecision.LOCAL_WITH_TOOLS if tools_needed else RoutingDecision.LOCAL_ONLY,
                'complexity': complexity,
                'confidence': confidence,
                'tools': tools_needed,
                'reasoning': f'Moderate task, OLMoE + {len(tools_needed)} tools'
            }
        
        elif complexity == TaskComplexity.UNCERTAIN:
            return {
                'strategy': RoutingDecision.SEARCH_FIRST,
                'complexity': complexity,
                'confidence': confidence,
                'tools': ['search'],
                'reasoning': 'Time-sensitive or factual, search first'
            }
        
        else:  # COMPLEX
            return {
                'strategy': RoutingDecision.DELEGATE,
                'complexity': complexity,
                'confidence': confidence,
                'tools': ['specialist_llm'],
                'reasoning': 'Complex task, needs specialist model'
            }
    
    def should_use_olmoe(self, routing: Dict) -> bool:
        """Check if OLMoE should be used"""
        return routing['strategy'] in [
            RoutingDecision.LOCAL_ONLY,
            RoutingDecision.LOCAL_WITH_TOOLS,
            RoutingDecision.SEARCH_FIRST
        ]
    
    def estimate_tokens(self, query: str, routing: Dict) -> int:
        """Estimate token usage for budgeting"""
        base_tokens = len(query.split()) * 1.3  # Rough estimate
        
        if routing['strategy'] == RoutingDecision.LOCAL_ONLY:
            return int(base_tokens * 50)  # Response multiplier
        elif routing['strategy'] == RoutingDecision.LOCAL_WITH_TOOLS:
            return int(base_tokens * 80)  # Tools add overhead
        elif routing['strategy'] == RoutingDecision.SEARCH_FIRST:
            return int(base_tokens * 100)  # Search + synthesis
        else:
            return int(base_tokens * 200)  # Specialist is expensive
    
    def log_decision(self, query: str, routing: Dict, was_successful: bool):
        """Log routing decisions for analysis"""
        log_entry = {
            'query': query[:100],
            'strategy': routing['strategy'].value,
            'complexity': routing['complexity'].name,
            'confidence': routing['confidence'],
            'tools_used': routing['tools'],
            'success': was_successful
        }
        
        # In production, write to file/DB
        print(f"📊 Routing Log: {log_entry}")


def test_router():
    """Test ISL router with sample queries"""
    router = ISLRouter()
    
    test_queries = [
        "What is 2 + 2?",
        "Explain quantum computing",
        "Who won the latest election?",
        "Prove the Riemann hypothesis",
        "Write a Python function to reverse a string",
        "What's the weather today?",
        "Fix this grammar: Me and him goes to store",
    ]
    
    print("="*70)
    print("🧪 ISL Router Test")
    print("="*70)
    
    for query in test_queries:
        routing = router.route(query)
        tokens = router.estimate_tokens(query, routing)
        
        print(f"\n📝 Query: {query}")
        print(f"   Strategy: {routing['strategy'].value}")
        print(f"   Complexity: {routing['complexity'].name}")
        print(f"   Confidence: {routing['confidence']:.0%}")
        print(f"   Tools: {routing['tools'] or 'None'}")
        print(f"   Est. Tokens: ~{tokens}")
        print(f"   Reasoning: {routing['reasoning']}")


if __name__ == "__main__":
    test_router()