llm/manager.py

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import Any

from core.logging.logger import get_logger
from llm.providers.base_provider import BaseProvider, LLMProviderError, LLMRequest, LLMResponse

logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------

class LLMManagerError(Exception):
    """Base exception for LLMManager errors."""


class ConfigurationError(LLMManagerError):
    """Raised when LLMManager is misconfigured."""


class ProviderNotFoundError(LLMManagerError):
    """Raised when a requested provider is not registered."""


class ProviderAlreadyRegisteredError(LLMManagerError):
    """Raised when a provider with the same name is already registered."""


class AllProvidersFailedError(LLMManagerError):
    """Raised when all fallback providers have been exhausted."""

    def __init__(self, attempts: list[dict[str, Any]]) -> None:
        self.attempts = attempts
        provider_names = [a["provider"] for a in attempts]
        errors = [a["error"] for a in attempts]
        super().__init__(
            f"All providers failed: {provider_names}. Errors: {errors}"
        )


# ---------------------------------------------------------------------------
# Fallback Policy
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class FallbackPolicy:
    """Determines which providers to try and in what order.

    If ``providers`` is empty, the manager falls back to iterating over all
    registered providers in insertion order.

    Attributes:
        providers: Ordered list of provider names to try.
        max_retries_per_provider: How many times to retry each provider
            before moving to the next one.
        base_delay_seconds: Base delay for exponential backoff between
            retries (seconds).
        max_delay_seconds: Maximum delay between retries (seconds).
    """

    providers: list[str] = field(default_factory=list)
    max_retries_per_provider: int = 2
    base_delay_seconds: float = 1.0
    max_delay_seconds: float = 30.0


# ---------------------------------------------------------------------------
# LLMManager
# ---------------------------------------------------------------------------

class LLMManager:
    """Registry and dispatcher for LLM providers.

    Implements the Strategy pattern with retry and fallback:

    - ``generate()`` delegates to the specified provider (or default).
    - On ``LLMProviderError``, retries with exponential backoff.
    - When retries are exhausted, falls back to the next provider in the
      ``FallbackPolicy``.
    - Thread-safe provider registration via ``asyncio.Lock``.

    Usage::

        manager = LLMManager(default_provider="openai")
        manager.register_provider("openai", openai_provider)
        manager.register_provider("claude", claude_provider)

        response = await manager.generate(request, provider_name="openai")
    """

    def __init__(
        self,
        default_provider: str = "openai",
        fallback_policy: FallbackPolicy | None = None,
    ) -> None:
        self._providers: dict[str, BaseProvider] = {}
        self._default_provider: str = default_provider
        self._fallback_policy: FallbackPolicy = fallback_policy or FallbackPolicy()
        self._lock: asyncio.Lock = asyncio.Lock()

    async def register_provider(
        self, name: str, provider: BaseProvider, overwrite: bool = False
    ) -> None:
        """Register an LLM provider under the given name.

        Args:
            name: Unique provider identifier (e.g. "openai", "claude").
            provider: Provider instance implementing BaseProvider.
            overwrite: If True, silently replaces an existing provider.
                If False, raises ProviderAlreadyRegisteredError.

        Raises:
            ProviderAlreadyRegisteredError: If name is already registered
                and overwrite is False.
        """
        async with self._lock:
            if name in self._providers and not overwrite:
                raise ProviderAlreadyRegisteredError(
                    f"Provider '{name}' is already registered. "
                    f"Use overwrite=True to replace it."
                )

            self._providers[name] = provider

        logger.info(
            "provider_registered",
            provider_name=name,
            provider_class=provider.__class__.__name__,
        )

    def get_provider(self, name: str | None = None) -> BaseProvider:
        """Retrieve a registered provider by name.

        If name is None, returns the default provider.

        Args:
            name: Provider identifier, or None for the default.

        Returns:
            The registered BaseProvider instance.

        Raises:
            ProviderNotFoundError: If the requested provider is not found.
        """
        lookup_name = name or self._default_provider

        provider = self._providers.get(lookup_name)
        if provider is None:
            available = list(self._providers.keys())
            raise ProviderNotFoundError(
                f"Provider '{lookup_name}' not found. "
                f"Available providers: {available}"
            )

        return provider

    def _resolve_fallback_order(self, explicit_name: str | None = None) -> list[str]:
        """Determine the ordered list of provider names to attempt.

        If an explicit provider name is given, start with it and then
        fall back through the policy (or all registered providers).

        If no explicit name, start with the default provider.
        """
        policy_names = (
            list(self._fallback_policy.providers)
            if self._fallback_policy.providers
            else list(self._providers.keys())
        )

        primary = explicit_name or self._default_provider

        # Build ordered list: primary first, then policy order (excluding primary)
        order: list[str] = []
        if primary in self._providers:
            order.append(primary)

        for name in policy_names:
            if name not in order and name in self._providers:
                order.append(name)

        return order

    async def _attempt_with_retries(
        self,
        provider: BaseProvider,
        request: LLMRequest,
        max_retries: int,
        base_delay: float,
        max_delay: float,
    ) -> LLMResponse:
        """Call provider.generate() with exponential backoff retries.

        Args:
            provider: The provider instance.
            request: The LLM request.
            max_retries: Maximum number of retries after the first attempt.
            base_delay: Base delay for exponential backoff (seconds).
            max_delay: Maximum delay between retries (seconds).

        Returns:
            LLMResponse on success.

        Raises:
            LLMProviderError: If all retries are exhausted.
        """
        last_error: LLMProviderError | None = None
        total_attempts = max_retries + 1

        for attempt in range(total_attempts):
            try:
                response = await provider.generate(request)

                if attempt > 0:
                    logger.info(
                        "llm_generate_retry_succeeded",
                        provider=provider.name,
                        attempt=attempt + 1,
                        total_attempts=total_attempts,
                    )

                return response

            except LLMProviderError as exc:
                last_error = exc
                if attempt < max_retries:
                    delay = min(base_delay * (2 ** attempt), max_delay)
                    logger.warning(
                        "llm_generate_retry",
                        provider=provider.name,
                        attempt=attempt + 1,
                        max_retries=max_retries + 1,
                        retry_in_seconds=delay,
                        error=str(exc),
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(
                        "llm_generate_exhausted_retries",
                        provider=provider.name,
                        total_attempts=total_attempts,
                        error=str(exc),
                    )

        # Should never reach here, but for type safety
        raise last_error or LLMProviderError(  # type: ignore[misc]
            "Unexpected retry loop exit",
            provider_name=provider.name,
        )

    async def generate(
        self,
        request: LLMRequest,
        provider_name: str | None = None,
    ) -> LLMResponse:
        """Generate a response using the specified or default provider.

        Implements retry with exponential backoff per provider, then
        fallback to the next provider in the FallbackPolicy order.

        Args:
            request: The LLM request to send.
            provider_name: Optional provider to use; falls back to default.

        Returns:
            The LLMResponse from the provider.

        Raises:
            ProviderNotFoundError: If no providers are registered
                at all.
            AllProvidersFailedError: If all providers fail after retries.
        """
        fallback_order = self._resolve_fallback_order(provider_name)

        if not fallback_order:
            raise ProviderNotFoundError(
                f"No providers registered. "
                f"Requested: '{provider_name or self._default_provider}'"
            )

        policy = self._fallback_policy
        attempts: list[dict[str, Any]] = []

        for name in fallback_order:
            provider = self._providers.get(name)
            if provider is None:
                logger.debug(
                    "llm_generate_skip_missing_provider",
                    provider_name=name,
                )
                continue

            logger.debug(
                "llm_generate_attempt",
                provider=provider.name,
                model=request.model,
                messages_count=len(request.messages),
                attempt=len(attempts) + 1,
                total_providers=len(fallback_order),
            )

            try:
                response = await self._attempt_with_retries(
                    provider=provider,
                    request=request,
                    max_retries=policy.max_retries_per_provider,
                    base_delay=policy.base_delay_seconds,
                    max_delay=policy.max_delay_seconds,
                )

                logger.debug(
                    "llm_generate_complete",
                    provider=provider.name,
                    model=response.model,
                    content_length=len(response.content),
                )

                return response

            except LLMProviderError as exc:
                attempts.append({
                    "provider": name,
                    "error": str(exc),
                })
                logger.warning(
                    "llm_generate_provider_failed",
                    provider=name,
                    error=str(exc),
                    remaining_providers=[
                        n for n in fallback_order
                        if n not in [a["provider"] for a in attempts]
                    ],
                )

        raise AllProvidersFailedError(attempts)

    def list_providers(self) -> list[str]:
        """Return a list of all registered provider names."""
        return list(self._providers.keys())

    @property
    def default_provider(self) -> str:
        """Return the name of the default provider."""
        return self._default_provider

    @default_provider.setter
    def default_provider(self, name: str, strict: bool = False) -> None:
        """Set the default provider name.

        Args:
            name: Provider name to set as default.
            strict: If True, raises ValueError when the provider is not
                registered. If False, logs a warning but allows the assignment.

        Raises:
            ValueError: In strict mode, when provider is not registered.
        """
        if name not in self._providers:
            if strict:
                raise ValueError(
                    f"Cannot set default_provider to '{name}': "
                    f"provider not registered. "
                    f"Available: {self.list_providers()}"
                )
            logger.warning(
                "default_provider_set_to_unregistered",
                provider_name=name,
                available=list(self._providers.keys()),
            )
        self._default_provider = name

    def has_providers(self) -> bool:
        """Return True if at least one provider is registered."""
        return len(self._providers) > 0

    @property
    def fallback_policy(self) -> FallbackPolicy:
        """Return the current fallback policy."""
        return self._fallback_policy

    @fallback_policy.setter
    def fallback_policy(self, policy: FallbackPolicy) -> None:
        """Set a new fallback policy."""
        self._fallback_policy = policy

    def __repr__(self) -> str:
        return (
            f"<LLMManager providers={self.list_providers()} "
            f"default='{self._default_provider}'>"
        )