#!/usr/bin/env python3
"""
Task routing system for local-first AI fallbacks.
Issue #483: [AUDIT][RISK] Maintain local-first fallbacks for all cloud AI
"""
import json
import time
import argparse
from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from enum import Enum
import requests

class TaskType(Enum):
    """Types of tasks that can be routed."""
    CODE_GENERATION = "code_generation"
    WEB_SEARCH = "web_search"
    DOCUMENT_ANALYSIS = "document_analysis"
    CREATIVE_WRITING = "creative_writing"
    DATA_ANALYSIS = "data_analysis"
    QUESTION_ANSWERING = "question_answering"
    SUMMARIZATION = "summarization"

class Priority(Enum):
    """Routing priority."""
    LOCAL_FIRST = "local-first"
    QUALITY_FIRST = "quality-first"
    BALANCED = "balanced"
    COST_FIRST = "cost-first"

@dataclass
class TaskResult:
    """Result from a task execution."""
    provider: str
    model: str
    result: str
    quality_score: float
    execution_time: float
    cost: float
    timestamp: str

class LocalModel:
    """Interface to local Ollama models."""
    
    def __init__(self, name: str, endpoint: str = "http://localhost:11434"):
        self.name = name
        self.endpoint = endpoint
        self.available = self._check_availability()
    
    def _check_availability(self) -> bool:
        """Check if model is available."""
        try:
            response = requests.get(f"{self.endpoint}/api/tags", timeout=5)
            if response.status_code == 200:
                models = [m["name"] for m in response.json().get("models", [])]
                return self.name in models
        except:
            pass
        return False
    
    def execute(self, prompt: str, max_tokens: int = 100) -> Optional[str]:
        """Execute a prompt on the local model."""
        if not self.available:
            return None
        
        try:
            payload = {
                "model": self.name,
                "prompt": prompt,
                "stream": False,
                "options": {"num_predict": max_tokens}
            }
            
            response = requests.post(
                f"{self.endpoint}/api/generate",
                json=payload,
                timeout=30
            )
            
            if response.status_code == 200:
                return response.json().get("response", "")
        except Exception as e:
            print(f"Error executing local model {self.name}: {e}")
        
        return None

class ExternalService:
    """Interface to external AI services."""
    
    def __init__(self, name: str, api_key: Optional[str] = None):
        self.name = name
        self.api_key = api_key
        self.available = self._check_availability()
    
    def _check_availability(self) -> bool:
        """Check if service is available."""
        # Simplified check - in reality would test actual API
        if self.name == "perplexity":
            return bool(self.api_key)
        return False
    
    def execute(self, prompt: str) -> Optional[str]:
        """Execute a prompt on the external service."""
        if not self.available:
            return None
        
        # Simplified implementation
        # In reality, would call actual API
        return f"External {self.name} response to: {prompt[:50]}..."

class TaskRouter:
    """Routes tasks to appropriate providers based on priority."""
    
    def __init__(self):
        self.local_models = {
            "hermes4": LocalModel("hermes4"),
            "llama3-8b": LocalModel("llama3-8b"),
            "mistral-7b": LocalModel("mistral-7b")
        }
        
        self.external_services = {
            "perplexity": ExternalService("perplexity")
        }
        
        # Task capability matrix
        self.capabilities = {
            TaskType.CODE_GENERATION: {
                "local": ["hermes4", "llama3-8b"],
                "external": [],
                "quality_local": 0.8,
                "quality_external": 0.9
            },
            TaskType.WEB_SEARCH: {
                "local": [],
                "external": ["perplexity"],
                "quality_local": 0.3,
                "quality_external": 0.95
            },
            TaskType.DOCUMENT_ANALYSIS: {
                "local": ["hermes4", "llama3-8b", "mistral-7b"],
                "external": [],
                "quality_local": 0.85,
                "quality_external": 0.9
            },
            TaskType.CREATIVE_WRITING: {
                "local": ["hermes4", "mistral-7b"],
                "external": [],
                "quality_local": 0.75,
                "quality_external": 0.85
            }
        }
    
    def route_task(self, task_type: TaskType, prompt: str, priority: Priority = Priority.BALANCED) -> Optional[TaskResult]:
        """Route a task based on priority."""
        print(f"Routing {task_type.value} with {priority.value} priority...")
        
        if priority == Priority.LOCAL_FIRST:
            return self._route_local_first(task_type, prompt)
        elif priority == Priority.QUALITY_FIRST:
            return self._route_quality_first(task_type, prompt)
        elif priority == Priority.COST_FIRST:
            return self._route_cost_first(task_type, prompt)
        else:  # BALANCED
            return self._route_balanced(task_type, prompt)
    
    def _route_local_first(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
        """Try local models first, fallback to external."""
        print("  Trying local models first...")
        
        # Try local models
        for model_name in self.capabilities.get(task_type, {}).get("local", []):
            model = self.local_models.get(model_name)
            if model and model.available:
                print(f"    Trying {model_name}...")
                start_time = time.time()
                result = model.execute(prompt)
                exec_time = time.time() - start_time
                
                if result:
                    return TaskResult(
                        provider="local",
                        model=model_name,
                        result=result,
                        quality_score=self.capabilities[task_type]["quality_local"],
                        execution_time=exec_time,
                        cost=0.0,
                        timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
                    )
        
        # Fallback to external
        print("  Local failed, trying external...")
        for service_name in self.capabilities.get(task_type, {}).get("external", []):
            service = self.external_services.get(service_name)
            if service and service.available:
                start_time = time.time()
                result = service.execute(prompt)
                exec_time = time.time() - start_time
                
                if result:
                    return TaskResult(
                        provider="external",
                        model=service_name,
                        result=result,
                        quality_score=self.capabilities[task_type]["quality_external"],
                        execution_time=exec_time,
                        cost=0.01,  # Estimated cost
                        timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
                    )
        
        return None
    
    def _route_quality_first(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
        """Choose provider with highest quality."""
        print("  Choosing highest quality provider...")
        
        best_result = None
        best_quality = 0
        
        # Check external first (usually higher quality)
        for service_name in self.capabilities.get(task_type, {}).get("external", []):
            service = self.external_services.get(service_name)
            if service and service.available:
                quality = self.capabilities[task_type]["quality_external"]
                if quality > best_quality:
                    start_time = time.time()
                    result = service.execute(prompt)
                    exec_time = time.time() - start_time
                    
                    if result:
                        best_quality = quality
                        best_result = TaskResult(
                            provider="external",
                            model=service_name,
                            result=result,
                            quality_score=quality,
                            execution_time=exec_time,
                            cost=0.01,
                            timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
                        )
        
        # Check local models
        for model_name in self.capabilities.get(task_type, {}).get("local", []):
            model = self.local_models.get(model_name)
            if model and model.available:
                quality = self.capabilities[task_type]["quality_local"]
                if quality > best_quality:
                    start_time = time.time()
                    result = model.execute(prompt)
                    exec_time = time.time() - start_time
                    
                    if result:
                        best_quality = quality
                        best_result = TaskResult(
                            provider="local",
                            model=model_name,
                            result=result,
                            quality_score=quality,
                            execution_time=exec_time,
                            cost=0.0,
                            timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
                        )
        
        return best_result
    
    def _route_cost_first(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
        """Choose cheapest provider."""
        print("  Choosing cheapest provider...")
        
        # Try local first (free)
        for model_name in self.capabilities.get(task_type, {}).get("local", []):
            model = self.local_models.get(model_name)
            if model and model.available:
                start_time = time.time()
                result = model.execute(prompt)
                exec_time = time.time() - start_time
                
                if result:
                    return TaskResult(
                        provider="local",
                        model=model_name,
                        result=result,
                        quality_score=self.capabilities[task_type]["quality_local"],
                        execution_time=exec_time,
                        cost=0.0,
                        timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
                    )
        
        # Fallback to external (paid)
        for service_name in self.capabilities.get(task_type, {}).get("external", []):
            service = self.external_services.get(service_name)
            if service and service.available:
                start_time = time.time()
                result = service.execute(prompt)
                exec_time = time.time() - start_time
                
                if result:
                    return TaskResult(
                        provider="external",
                        model=service_name,
                        result=result,
                        quality_score=self.capabilities[task_type]["quality_external"],
                        execution_time=exec_time,
                        cost=0.01,
                        timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
                    )
        
        return None
    
    def _route_balanced(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
        """Balance quality and cost."""
        print("  Balancing quality and cost...")
        
        # Simple heuristic: if local quality is within 10% of external, use local
        local_quality = self.capabilities.get(task_type, {}).get("quality_local", 0)
        external_quality = self.capabilities.get(task_type, {}).get("quality_external", 0)
        
        if local_quality >= external_quality * 0.9:
            # Local quality is good enough
            return self._route_local_first(task_type, prompt)
        else:
            # External quality is significantly better
            return self._route_quality_first(task_type, prompt)

def main():
    parser = argparse.ArgumentParser(description="Route tasks to local or external AI providers")
    parser.add_argument("--task", required=True, choices=[t.value for t in TaskType],
                       help="Task type")
    parser.add_argument("--prompt", required=True, help="Task prompt")
    parser.add_argument("--priority", default="balanced", 
                       choices=[p.value for p in Priority],
                       help="Routing priority")
    parser.add_argument("--output", help="Output file for results")
    
    args = parser.parse_args()
    
    # Create router
    router = TaskRouter()
    
    # Route task
    task_type = TaskType(args.task)
    priority = Priority(args.priority)
    
    result = router.route_task(task_type, args.prompt, priority)
    
    if result:
        print(f"\nTask routed to: {result.provider} ({result.model})")
        print(f"Quality score: {result.quality_score:.2f}")
        print(f"Execution time: {result.execution_time:.2f}s")
        print(f"Cost: ${result.cost:.4f}")
        print(f"Result preview: {result.result[:100]}...")
        
        # Save results if requested
        if args.output:
            with open(args.output, 'w') as f:
                json.dump({
                    "task_type": task_type.value,
                    "prompt": args.prompt,
                    "priority": priority.value,
                    "result": {
                        "provider": result.provider,
                        "model": result.model,
                        "quality_score": result.quality_score,
                        "execution_time": result.execution_time,
                        "cost": result.cost,
                        "timestamp": result.timestamp,
                        "result": result.result
                    }
                }, f, indent=2)
            print(f"Results saved to {args.output}")
    else:
        print("Failed to route task - no available providers")
        return 1
    
    return 0

if __name__ == "__main__":
    import sys
    sys.exit(main())