"""
Core router implementation for InferenceKit.
Handles routing between different providers and models.
"""

import asyncio
from typing import Dict, Any, List, Optional, Union
from datetime import datetime
from uuid import UUID, uuid4
from enum import Enum

from ..types import (
    ModelType,
    ProviderType,
    Request,
    Response,
    ModelConfig,
    BatchRequest,
    BatchResponse,
)
from ..exceptions import (
    ModelNotFoundError,
    ProviderError,
    InvalidRequestError,
    ConfigurationError,
)
from ..core.model import Model
from ..core.cache import SemanticCache
from ..core.streaming import StreamingManager
from ..monitoring.metrics import Metrics
from ..monitoring.logger import Logger
from ..monitoring.tracer import Tracer


class Router:
    """Main router class for handling model routing and orchestration."""

    class RoutingStrategy(str, Enum):
        ROUND_ROBIN = "round_robin"
        LOWEST_LATENCY = "lowest_latency"
        LOWEST_COST = "lowest_cost"
        HIGHEST_QUALITY = "highest_quality"
        FAILOVER = "failover"

    def __init__(
        self,
        models: List[ModelConfig],
        routing_strategy: RoutingStrategy = RoutingStrategy.ROUND_ROBIN,
        cache_config: Optional[Dict[str, Any]] = None,
        monitoring_config: Optional[Dict[str, Any]] = None,
        streaming_config: Optional[Dict[str, Any]] = None
    ):
        self.models = models
        self.routing_strategy = routing_strategy
        self.cache_config = cache_config or {}
        self.monitoring_config = monitoring_config or {}
        self.streaming_config = streaming_config or {}

        self._model_instances: Dict[str, Model] = {}
        self._current_model_index = 0
        self._metrics_collector = MetricsCollector(**self.monitoring_config.get("metrics", {}))
        self._logger = Logger(**self.monitoring_config.get("logger", {}))
        self._tracer = Tracer(**self.monitoring_config.get("tracer", {}))
        self._streaming_manager = StreamingManager(**self.streaming_config)

        # Initialize routing state
        self._routing_stats = {
            "requests": 0,
            "successes": 0,
            "failures": 0,
            "latencies": [],
            "costs": [],
        }

    async def initialize(self) -> None:
        """Initialize all models and components."""
        self._logger.info("Initializing router with %d models", len(self.models))
        self._tracer.start_trace("router_initialization")

        for model_config in self.models:
            try:
                model = Model(model_config)
                await model.initialize()
                self._model_instances[model_config.model_name] = model
                self._logger.info("Initialized model: %s", model_config.model_name)
            except Exception as e:
                self._logger.error("Failed to initialize model %s: %s", model_config.model_name, str(e))
                raise ConfigurationError(f"Failed to initialize model {model_config.model_name}")

        self._tracer.end_trace("router_initialization")

    async def shutdown(self) -> None:
        """Shutdown all models and components."""
        self._logger.info("Shutting down router")
        self._tracer.start_trace("router_shutdown")

        for model_name, model in self._model_instances.items():
            try:
                await model.shutdown()
                self._logger.info("Shutdown model: %s", model_name)
            except Exception as e:
                self._logger.error("Failed to shutdown model %s: %s", model_name, str(e))

        self._streaming_manager.shutdown()
        self._metrics_collector.shutdown()
        self._tracer.shutdown()

        self._tracer.end_trace("router_shutdown")

    async def _route_request(self, request: Request) -> Model:
        """Route a request to the appropriate model."""
        if not self._model_instances:
            raise ModelNotFoundError("No models available for routing")

        # Get available models for the request type
        available_models = [
            (name, model) for name, model in self._model_instances.items()
            if model.config.model_name in request.model
        ]

        if not available_models:
            raise ModelNotFoundError(f"No models found for model name: {request.model}")

        # Apply routing strategy
        if self.routing_strategy == self.RoutingStrategy.ROUND_ROBIN:
            return await self._route_round_robin(request, available_models)
        elif self.routing_strategy == self.RoutingStrategy.LOWEST_LATENCY:
            return await self._route_lowest_latency(request, available_models)
        elif self.routing_strategy == self.RoutingStrategy.LOWEST_COST:
            return await self._route_lowest_cost(request, available_models)
        elif self.routingStrategy == self.RoutingStrategy.HIGHEST_QUALITY:
            return await self._route_highest_quality(request, available_models)
        elif self.routing_strategy == self.RoutingStrategy.FAILOVER:
            return await self._route_failover(request, available_models)
        else:
            raise InvalidRequestError(f"Unknown routing strategy: {self.routing_strategy}")

    async def _route_round_robin(self, request: Request, available_models: List[tuple]) -> Model:
        """Route using round-robin strategy."""
        model_name, model = available_models[self._current_model_index % len(available_models)]
        self._current_model_index += 1
        return model

    async def _route_lowest_latency(self, request: Request, available_models: List[tuple]) -> Model:
        """Route using lowest latency strategy."""
        # In a real implementation, this would use historical latency data
        # For now, we'll just use round-robin as a placeholder
        return await self._route_round_robin(request, available_models)

    async def _route_lowest_cost(self, request: Request, available_models: List[tuple]) -> Model:
        """Route using lowest cost strategy."""
        # In a real implementation, this would use cost data
        # For now, we'll just use round-robin as a placeholder
        return await self._route_round_robin(request, available_models)

    async def _route_highest_quality(self, request: Request, available_models: List[tuple]) -> Model:
        """Route using highest quality strategy."""
        # In a real implementation, this would use quality metrics
        # For now, we'll just use round-robin as a placeholder
        return await self._route_round_robin(request, available_models)

    async def _route_failover(self, request: Request, available_models: List[tuple]) -> Model:
        """Route using failover strategy."""
        # Try models in order until one succeeds
        for model_name, model in available_models:
            try:
                # Test if model is healthy
                if await model.get_provider_limits():
                    return model
            except Exception:
                continue

        raise ModelNotFoundError("No healthy models available for failover")

    async def generate(
        self,
        prompt: str,
        model: Optional[str] = None,
        provider: Optional[ProviderType] = None,
        **kwargs
    ) -> Response:
        """Generate text using the router."""
        request = Request(
            model=model or self.models[0].model_name,
            provider=provider or self.models[0].provider,
            type=ModelType.TEXT_GENERATION,
            prompt=prompt,
            parameters=kwargs,
        )

        return await self._process_request(request)

    async def chat(
        self,
        messages: List[Dict[str, Any]],
        model: Optional[str] = None,
        provider: Optional[ProviderType] = None,
        **kwargs
    ) -> Response:
        """Chat using the router."""
        request = Request(
            model=model or self.models[0].model_name,
            provider=provider or self.models[0].provider,
            type=ModelType.CHAT,
            prompt="\n".join([
                f"System: {msg.get('role', 'system')}: {msg.get('content', '')}"
                for msg in messages
            ]),
            parameters=kwargs,
        )

        return await self._process_request(request)

    async def embed(
        self,
        text: str,
        model: Optional[str] = None,
        provider: Optional[ProviderType] = None,
        **kwargs
    ) -> Response:
        """Generate embeddings using the router."""
        request = Request(
            model=model or self.models[0].model_name,
            provider=provider or self.models[0].provider,
            type=ModelType.EMBEDDING,
            prompt=text,
            parameters=kwargs,
        )

        return await self._process_request(request)

    async def batch_generate(
        self,
        prompts: List[str],
        model: Optional[str] = None,
        provider: Optional[ProviderType] = None,
        batch_size: int = 10,
        max_concurrency: int = 5,
        **kwargs
    ) -> BatchResponse:
        """Generate multiple responses in batches."""
        batch_request = BatchRequest(
            requests=[
                Request(
                    model=model or self.models[0].model_name,
                    provider=provider or self.models[0].provider,
                    type=ModelType.TEXT_GENERATION,
                    prompt=prompt,
                    parameters=kwargs,
                )
                for prompt in prompts
            ],
            batch_size=batch_size,
            max_concurrency=max_concurrency,
        )

        return await self._process_batch_request(batch_request)

    async def _process_request(self, request: Request) -> Response:
        """Process a single request through the router."""
        self._metrics_collector.increment_requests()
        self._logger.info("Processing request: %s", request.id)
        self._tracer.start_trace("request_processing", request_id=str(request.id))

        try:
            # Route to appropriate model
            model = await self._route_request(request)

            # Process request
            start_time = datetime.utcnow()
            response = await model.generate(request.prompt, **request.parameters)
            end_time = datetime.utcnow()

            # Record metrics
            latency = (end_time - start_time).total_seconds()
            self._routing_stats["requests"] += 1
            self._routing_stats["successes"] += 1
            self._routing_stats["latencies"].append(latency)

            # Log and trace
            self._logger.info("Request %s completed successfully in %.2f seconds", request.id, latency)
            self._tracer.log_event("request_completed", {
                "latency": latency,
                "model": model.config.model_name
            })

            # Update metrics
            self._metrics_collector.record_latency(latency)
            self._metrics_collector.record_success()

            return response

        except Exception as e:
            self._metrics_collector.increment_failures()
            self._routing_stats["failures"] += 1
            self._logger.error("Request %s failed: %s", request.id, str(e))
            self._tracer.log_error("request_failed", str(e))
            raise

        finally:
            self._tracer.end_trace("request_processing")

    async def _process_batch_request(self, batch_request: BatchRequest) -> BatchResponse:
        """Process a batch request through the router."""
        self._metrics_collector.increment_requests(len(batch_request.requests))
        self._logger.info("Processing batch request with %d items", len(batch_request.requests))
        self._tracer.start_trace("batch_request_processing")

        results = []
        failed_requests = []
        start_time = datetime.utcnow()

        try:
            for i in range(0, len(batch_request.requests), batch_request.batch_size):
                batch = batch_request.requests[i:i + batch_request.batch_size]

                tasks = []
                semaphore = asyncio.Semaphore(batch_request.max_concurrency)

                async def bounded_request(request: Request):
                    async with semaphore:
                        try:
                            return await self._process_request(request)
                        except Exception:
                            failed_requests.append(request)
                            return None

                for request in batch:
                    tasks.append(bounded_request(request))

                batch_results = await asyncio.gather(*tasks)
                results.extend([r for r in batch_results if r is not None])

            total_time = (datetime.utcnow() - start_time).total_seconds()

            self._logger.info("Batch request completed with %d successes and %d failures",
                           len(results), len(failed_requests))

            return BatchResponse(
                results=results,
                failed_requests=failed_requests,
                total_time=total_time,
            )

        except Exception as e:
            self._metrics_collector.increment_failures()
            self._logger.error("Batch request failed: %s", str(e))
            self._tracer.log_error("batch_request_failed", str(e))
            raise

        finally:
            self._tracer.end_trace("batch_request_processing")

    async def get_model_info(self, model_name: str) -> Dict[str, Any]:
        """Get information about a specific model."""
        if model_name not in self._model_instances:
            raise ModelNotFoundError(f"Model {model_name} not found")

        return await self._model_instances[model_name].get_model_info()

    async def list_models(self) -> List[str]:
        """List all available models."""
        return list(self._model_instances.keys())

    async def get_stats(self) -> Dict[str, Any]:
        """Get router statistics."""
        return {
            "total_requests": self._routing_stats["requests"],
            "successful_requests": self._routing_stats["successes"],
            "failed_requests": self._routing_stats["failures"],
            "average_latency": sum(self._routing_stats["latencies"]) / len(self._routing_stats["latencies"]) if self._routing_stats["latencies"] else 0,
            "total_costs": sum(self._routing_stats["costs"]) if self._routing_stats["costs"] else 0,
            "monitoring": self._metrics_collector.get_metrics(),
        }

    async def get_health(self) -> Dict[str, Any]:
        """Get health status of all models."""
        health_status = {}
        for model_name, model in self._model_instances.items():
            try:
                health_status[model_name] = {
                    "healthy": await model.get_provider_limits(),
                    "model_info": await model.get_model_info()
                }
            except Exception as e:
                health_status[model_name] = {
                    "healthy": False,
                    "error": str(e)
                }
        return health_status

    async def update_routing_strategy(self, strategy: RoutingStrategy) -> None:
        """Update the routing strategy."""
        self.routing_strategy = strategy
        self._logger.info("Updated routing strategy to: %s", strategy.value)

    def get_routing_strategy(self) -> RoutingStrategy:
        """Get the current routing strategy."""
        return self.routing_strategy

    async def add_model(self, model_config: ModelConfig) -> None:
        """Add a new model to the router."""
        if model_config.model_name in self._model_instances:
            raise ConfigurationError(f"Model {model_config.model_name} already exists")

        model = Model(model_config)
        await model.initialize()
        self._model_instances[model_config.model_name] = model
        self._logger.info("Added new model: %s", model_config.model_name)

    async def remove_model(self, model_name: str) -> None:
        """Remove a model from the router."""
        if model_name not in self._model_instances:
            raise ModelNotFoundError(f"Model {model_name} not found")

        model = self._model_instances.pop(model_name)
        await model.shutdown()
        self._logger.info("Removed model: %s", model_name)

    async def get_cache_stats(self) -> Dict[str, Any]:
        """Get cache statistics across all models."""
        cache_stats = {}
        for model_name, model in self._model_instances.items():
            cache_stats[model_name] = await model.get_cache_stats()
        return cache_stats

    async def clear_cache(self) -> None:
        """Clear cache across all models."""
        for model_name, model in self._model_instances.items():
            await model.clear_cache()
        self._logger.info("Cleared cache across all models")