"""Model router with automatic fallback support."""

from typing import Any, AsyncIterator

from openai import (
    APIConnectionError,
    APIError,
    APIStatusError,
    APITimeoutError,
    InternalServerError,
    RateLimitError,
)
from openai.types.chat import ChatCompletion

from swarm.config.logging import get_structured_logger

from .factory import ProviderFactory
from .types import CompletionRequest

logger = get_structured_logger(__name__)


class ModelRouter:
    """Routes requests to appropriate providers with automatic fallback.

    Not a ModelProvider itself, but orchestrates multiple providers based on
    model prefixes and handles retry logic across different models."""

    def __init__(
        self,
        factory: ProviderFactory,
    ):
        """Initialize router.

        Args:
            factory: Provider factory used to resolve model prefixes
        """
        self.factory = factory

    def create_completion(self, request: CompletionRequest) -> ChatCompletion:
        """Create completion with automatic fallback on failure.

        Tries each model in order until one succeeds.
        """
        errors = []

        for i, model in enumerate(request.model):
            try:
                # Get provider for this model
                provider = self.factory.get_provider_for_model(model)

                # Pass request as-is, provider will extract model name
                return provider.create_completion(request, model)

            # Only catch retryable errors - others fail immediately
            except (
                APIConnectionError,  # Network/connection issues
                APITimeoutError,  # Request timeout
                InternalServerError,  # 500 errors
                APIStatusError,  # 503 errors
                APIError,  # Generic 500 errors
                RateLimitError,  # Rate limiting
            ) as e:
                errors.append((model, str(e)))
                match i:
                    case _ if i < len(request.model) - 1:
                        logger.debug(
                            f"Model {model} failed, trying next",
                            extra={"error": str(e)},
                        )
                        continue
                    case _:
                        logger.error(
                            "All models failed",
                            extra={"errors": errors},
                        )
                        raise RuntimeError(f"All models failed: {errors}") from e

        # Should never reach here due to raise above
        raise RuntimeError(f"All models failed: {errors}")

    async def create_completion_async(
        self, request: CompletionRequest
    ) -> ChatCompletion:
        """Create completion with automatic fallback (async).

        Tries each model in order until one succeeds.
        """
        errors = []

        for i, model in enumerate(request.model):
            try:
                # Get provider for this model
                provider = self.factory.get_provider_for_model(model)

                # Try completion
                return await provider.create_completion_async(request, model)

            # Only catch retryable errors - others fail immediately
            except (
                APIConnectionError,  # Network/connection issues
                APITimeoutError,  # Request timeout
                InternalServerError,  # 500 errors
                APIStatusError,  # 503 errors
                RateLimitError,  # Rate limiting
            ) as e:
                errors.append((model, str(e)))
                match i:
                    case _ if i < len(request.model) - 1:
                        logger.debug(
                            f"Model {model} failed, trying next",
                            extra={"error": str(e)},
                        )
                        continue
                    case _:
                        logger.error(
                            "All models failed",
                            extra={"errors": errors},
                        )
                        raise RuntimeError(f"All models failed: {errors}") from e

        # Should never reach here due to raise above
        raise RuntimeError(f"All models failed: {errors}")

    async def stream_completion(self, request: CompletionRequest) -> AsyncIterator[Any]:
        """Stream completion with automatic fallback.

        Tries each model in order until one succeeds.
        """
        errors = []

        for i, model in enumerate(request.model):
            try:
                # Get provider for this model
                provider = self.factory.get_provider_for_model(model)

                # Stream from this provider
                async for chunk in provider.stream_completion(request, model):
                    yield chunk

                # Success - streaming completed
                return

            # Only catch retryable errors - others fail immediately
            except (
                APIConnectionError,  # Network/connection issues
                APITimeoutError,  # Request timeout
                InternalServerError,  # 500 errors
                APIStatusError,  # 503 errors
                APIError,  # Generic 500 errors
                RateLimitError,  # Rate limiting
            ) as e:
                errors.append((model, str(e)))
                match i:
                    case _ if i < len(request.model) - 1:
                        logger.debug(
                            f"Model {model} failed, trying next",
                            extra={"error": str(e)},
                        )
                        continue
                    case _:
                        logger.error(
                            "All models failed to stream",
                            extra={"errors": errors},
                        )
                        raise RuntimeError(f"All models failed: {errors}") from e
