"""Fallback provider that switches from local to cloud on context overflow."""

from __future__ import annotations

from collections.abc import AsyncIterator
from contextlib import suppress
from typing import Any

from oclawma.providers.base import (
    BaseProvider,
    CompletionError,
    CompletionRequest,
    CompletionResponse,
    ProviderError,
)


class ContextOverflowError(CompletionError):
    """Raised when the context size exceeds the local model's limit."""

    def __init__(self, message: str = "Context too large for local model") -> None:
        super().__init__(message)


class FallbackProvider(BaseProvider):
    """Provider that falls back from local to cloud on context overflow.

    This provider wraps a primary (local) provider and a fallback (cloud) provider.
    When the primary provider fails due to context size limitations, it automatically
    switches to the fallback provider.

    Attributes:
        primary: The primary provider (typically local like Ollama).
        fallback: The fallback provider (typically cloud like Kimi).
        context_threshold: Token threshold before considering fallback.
    """

    # Typical context limits for local models
    DEFAULT_LOCAL_CONTEXT_LIMIT = 8192
    # Threshold at which to trigger fallback (80% of limit)
    FALLBACK_THRESHOLD_PERCENT = 0.8

    def __init__(
        self,
        primary: BaseProvider,
        fallback: BaseProvider,
        context_threshold: int | None = None,
    ) -> None:
        """Initialize the fallback provider.

        Args:
            primary: The primary provider to use for most requests.
            fallback: The fallback provider for context overflow scenarios.
            context_threshold: Token threshold to trigger fallback.
                             Defaults to 80% of local model limit.
        """
        # Initialize with primary's config (we don't use base_url/api_key directly)
        super().__init__("fallback", None)
        self.primary = primary
        self.fallback = fallback
        self.context_threshold = context_threshold or int(
            self.DEFAULT_LOCAL_CONTEXT_LIMIT * self.FALLBACK_THRESHOLD_PERCENT
        )
        self._last_used_fallback = False
        self._fallback_reason: str | None = None

    @property
    def last_used_fallback(self) -> bool:
        """Whether the fallback provider was used for the last request."""
        return self._last_used_fallback

    @property
    def fallback_reason(self) -> str | None:
        """The reason for using the fallback provider, if applicable."""
        return self._fallback_reason

    def _should_use_fallback(self, request: CompletionRequest) -> tuple[bool, str | None]:
        """Determine if fallback should be used for this request.

        Args:
            request: The completion request to evaluate.

        Returns:
            Tuple of (should_fallback, reason).
        """
        # Count tokens in the request
        token_count = self.count_message_tokens(request.messages)

        if token_count > self.context_threshold:
            return (
                True,
                f"Context size ({token_count} tokens) exceeds threshold ({self.context_threshold})",
            )

        return False, None

    def count_tokens(self, text: str, model: str | None = None) -> int:
        """Estimate token count using primary provider's method.

        Args:
            text: The text to count tokens for.
            model: Optional model name.

        Returns:
            Estimated token count.
        """
        return self.primary.count_tokens(text, model)

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        """Generate a chat completion with automatic fallback.

        First attempts with the primary provider. If that fails due to
        context size, automatically retries with the fallback provider.

        Args:
            request: The completion request.

        Returns:
            The completion response.

        Raises:
            ProviderError: If both primary and fallback fail.
        """
        should_fallback, reason = self._should_use_fallback(request)

        if should_fallback:
            # Use fallback directly if context is too large
            self._last_used_fallback = True
            self._fallback_reason = reason
            try:
                response = await self.fallback.complete(request)
                # Annotate response to indicate fallback was used
                return response
            except ProviderError:
                raise

        # Try primary first
        self._last_used_fallback = False
        self._fallback_reason = None

        try:
            return await self.primary.complete(request)
        except CompletionError as e:
            # Check if it's a context-related error
            error_msg = str(e).lower()
            if any(
                keyword in error_msg
                for keyword in ["context", "tokens", "too large", "exceeds", "maximum"]
            ):
                # Fall back to cloud provider
                self._last_used_fallback = True
                self._fallback_reason = f"Primary failed: {e}"
                return await self.fallback.complete(request)
            raise

    async def stream_complete(
        self, request: CompletionRequest
    ) -> AsyncIterator[CompletionResponse]:
        """Generate a streaming chat completion with automatic fallback.

        First attempts with the primary provider. If that fails due to
        context size, automatically retries with the fallback provider.

        Args:
            request: The completion request.

        Yields:
            Completion response chunks.

        Raises:
            ProviderError: If both primary and fallback fail.
        """
        should_fallback, reason = self._should_use_fallback(request)

        if should_fallback:
            # Use fallback directly if context is too large
            self._last_used_fallback = True
            self._fallback_reason = reason
            async for chunk in self.fallback.stream_complete(request):
                yield chunk
            return

        # Try primary first
        self._last_used_fallback = False
        self._fallback_reason = None

        chunks_yielded = False
        try:
            async for chunk in self.primary.stream_complete(request):
                chunks_yielded = True
                yield chunk
        except CompletionError as e:
            # Check if it's a context-related error
            error_msg = str(e).lower()
            if not chunks_yielded and any(
                keyword in error_msg
                for keyword in ["context", "tokens", "too large", "exceeds", "maximum"]
            ):
                # Fall back to cloud provider only if no chunks were yielded
                self._last_used_fallback = True
                self._fallback_reason = f"Primary failed: {e}"
                async for chunk in self.fallback.stream_complete(request):
                    yield chunk
            else:
                raise

    async def list_models(self) -> list[str]:
        """List available models from both providers.

        Returns:
            Combined list of available model names with provider prefixes.
        """
        primary_models = []
        fallback_models = []

        with suppress(ProviderError):
            primary_models = await self.primary.list_models()

        with suppress(ProviderError):
            fallback_models = await self.fallback.list_models()

        # Prefix models with their source
        return [f"local:{m}" for m in primary_models] + [f"cloud:{m}" for m in fallback_models]

    async def health_check(self) -> dict[str, Any]:
        """Check the health of both providers.

        Returns:
            Health status information for both providers.
        """
        primary_health = await self.primary.health_check()
        fallback_health = await self.fallback.health_check()

        return {
            "status": (
                "healthy"
                if primary_health.get("status") == "healthy"
                or fallback_health.get("status") == "healthy"
                else "unhealthy"
            ),
            "primary": primary_health,
            "fallback": fallback_health,
            "context_threshold": self.context_threshold,
            "last_used_fallback": self._last_used_fallback,
            "fallback_reason": self._fallback_reason,
        }

    def get_fallback_info(self) -> dict[str, Any]:
        """Get information about fallback behavior.

        Returns:
            Dictionary with fallback configuration and status.
        """
        return {
            "context_threshold": self.context_threshold,
            "local_context_limit": self.DEFAULT_LOCAL_CONTEXT_LIMIT,
            "last_used_fallback": self._last_used_fallback,
            "fallback_reason": self._fallback_reason,
            "primary_provider": type(self.primary).__name__,
            "fallback_provider": type(self.fallback).__name__,
        }
