"""Main orchestrator with Model-First Chain of Responsibility."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator, Awaitable, Callable
from pathlib import Path
from typing import Any

from llm_rotator._types import Candidate, LLMResponse, RoutingContext, StreamChunk, Usage
from llm_rotator.backends import InMemoryBackend
from llm_rotator.circuit_breaker import CircuitBreaker
from llm_rotator.clients import AbstractLLMClient
from llm_rotator.config import ModelGroupConfig, ProviderConfig, RotatorConfig
from llm_rotator.exceptions import (
    AllAttemptsFailedError,
    LLMRotatorError,
    ServerError,
)
from llm_rotator.logging import RoutingTrace
from llm_rotator.quota import QuotaManager, QuotaWarning

QuotaWarningCallback = Callable[[QuotaWarning], Awaitable[None]]

_default_logger = logging.getLogger("llm_rotator")


class LLMRotator:
    """Fault-tolerant LLM provider rotation orchestrator."""

    def __init__(
        self,
        config: RotatorConfig,
        clients: dict[str, AbstractLLMClient] | None = None,
        backend: InMemoryBackend | None = None,
        logger: logging.Logger | None = None,
        on_quota_warning: QuotaWarningCallback | None = None,
        warning_threshold: float = 0.8,
    ) -> None:
        self._config = config
        self._backend = backend or InMemoryBackend()
        self._circuit_breaker = CircuitBreaker(self._backend)
        self._quota_manager = QuotaManager(
            self._backend, config.providers, warning_threshold=warning_threshold
        )
        self._clients = clients or {}
        self._hooks: list[Any] = []
        self._logger = logger or _default_logger
        self._on_quota_warning = on_quota_warning

    @classmethod
    def from_json(
        cls,
        path: str | Path,
        clients: dict[str, AbstractLLMClient] | None = None,
        backend: InMemoryBackend | None = None,
        logger: logging.Logger | None = None,
        on_quota_warning: QuotaWarningCallback | None = None,
        warning_threshold: float = 0.8,
    ) -> LLMRotator:
        """Create a rotator from a JSON config file.

        Convenience wrapper around ``RotatorConfig.from_json()``.
        """
        config = RotatorConfig.from_json(path)
        return cls(
            config=config,
            clients=clients,
            backend=backend,
            logger=logger,
            on_quota_warning=on_quota_warning,
            warning_threshold=warning_threshold,
        )

    def add_hook(self, hook: Any) -> None:
        """Register a lifecycle hook."""
        self._hooks.append(hook)

    async def complete(
        self,
        messages: list[dict],
        routing: RoutingContext | None = None,
        **kwargs: Any,
    ) -> LLMResponse:
        """Execute a completion request with automatic rotation."""
        ctx = routing or RoutingContext()
        candidates = self._build_candidates(ctx)
        errors: list[LLMRotatorError] = []
        trace = RoutingTrace()

        for candidate in candidates:
            # 1. Circuit breaker check
            if not await self._circuit_breaker.is_candidate_available(candidate):
                trace.add_skip(candidate, "circuit breaker")
                continue

            # 2. Quota pre-check
            if not await self._quota_manager.pre_check(candidate):
                trace.add_skip(candidate, "quota exhausted")
                continue

            # 3. Hooks: before_request
            if not await self._run_before_hooks(ctx, candidate):
                trace.add_skip(candidate, "rejected by hook")
                continue

            # 4. Get client
            client = self._clients.get(candidate.client_type.value)
            if client is None:
                trace.add_skip(candidate, "no client")
                continue

            # 5. Retry loop
            provider_cfg = self._get_provider_config(candidate.provider_name)
            max_attempts = provider_cfg.server_error_retry.max_attempts if provider_cfg else 0
            delay = provider_cfg.server_error_retry.delay_seconds if provider_cfg else 0.0

            for attempt in range(max_attempts + 1):
                try:
                    response = await client.generate(
                        messages=messages,
                        model=candidate.model,
                        api_key=candidate.key_token,
                        base_url=candidate.base_url,
                        **kwargs,
                    )

                    # Record quota usage & check warnings
                    warnings = await self._quota_manager.record_usage(
                        candidate, response.usage
                    )
                    await self._emit_quota_warnings(warnings)

                    # Hooks: after_response
                    await self._run_after_hooks(ctx, candidate, response.usage)

                    trace.add_success(candidate, response.usage)
                    trace.emit(self._logger)
                    return response

                except ServerError as e:
                    errors.append(e)
                    await self._circuit_breaker.record_failure(candidate, e)

                    if attempt < max_attempts:
                        if delay > 0:
                            await asyncio.sleep(delay)
                        continue
                    trace.add_error(candidate, e)
                    break

                except LLMRotatorError as e:
                    errors.append(e)
                    await self._circuit_breaker.record_failure(candidate, e)
                    trace.add_error(candidate, e)
                    break

            # Notify fallback hooks
            if errors:
                await self._run_fallback_hooks(ctx, candidate, candidates, errors[-1])

        trace.add_all_failed()
        trace.emit(self._logger)
        raise AllAttemptsFailedError(
            f"All {len(errors)} attempts failed", errors or [LLMRotatorError("No candidates")]
        )

    async def stream(
        self,
        messages: list[dict],
        routing: RoutingContext | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        """Execute a streaming request with automatic rotation."""
        ctx = routing or RoutingContext()
        candidates = self._build_candidates(ctx)
        errors: list[LLMRotatorError] = []
        trace = RoutingTrace()

        for candidate in candidates:
            # Same pre-checks as complete()
            if not await self._circuit_breaker.is_candidate_available(candidate):
                trace.add_skip(candidate, "circuit breaker")
                continue
            if not await self._quota_manager.pre_check(candidate):
                trace.add_skip(candidate, "quota exhausted")
                continue
            if not await self._run_before_hooks(ctx, candidate):
                trace.add_skip(candidate, "rejected by hook")
                continue

            client = self._clients.get(candidate.client_type.value)
            if client is None:
                trace.add_skip(candidate, "no client")
                continue

            provider_cfg = self._get_provider_config(candidate.provider_name)
            max_attempts = provider_cfg.server_error_retry.max_attempts if provider_cfg else 0
            delay = provider_cfg.server_error_retry.delay_seconds if provider_cfg else 0.0

            for attempt in range(max_attempts + 1):
                try:
                    last_usage: Usage | None = None
                    async for chunk in client.stream(
                        messages=messages,
                        model=candidate.model,
                        api_key=candidate.key_token,
                        base_url=candidate.base_url,
                        **kwargs,
                    ):
                        if chunk.usage:
                            last_usage = chunk.usage
                        yield chunk

                    # Stream completed successfully
                    if last_usage:
                        warnings = await self._quota_manager.record_usage(
                            candidate, last_usage
                        )
                        await self._emit_quota_warnings(warnings)
                        await self._run_after_hooks(ctx, candidate, last_usage)
                        trace.add_success(candidate, last_usage)
                    else:
                        trace.add_success(
                            candidate,
                            Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
                        )
                    trace.emit(self._logger)
                    return

                except ServerError as e:
                    errors.append(e)
                    await self._circuit_breaker.record_failure(candidate, e)
                    if attempt < max_attempts:
                        if delay > 0:
                            await asyncio.sleep(delay)
                        continue
                    trace.add_error(candidate, e)
                    break

                except LLMRotatorError as e:
                    errors.append(e)
                    await self._circuit_breaker.record_failure(candidate, e)
                    trace.add_error(candidate, e)
                    break

            # Notify fallback hooks
            if errors:
                await self._run_fallback_hooks(ctx, candidate, candidates, errors[-1])

        trace.add_all_failed()
        trace.emit(self._logger)
        raise AllAttemptsFailedError(
            f"All {len(errors)} attempts failed",
            errors or [LLMRotatorError("No candidates")],
        )

    def _build_candidates(self, ctx: RoutingContext) -> list[Candidate]:
        """Build ordered candidate list: Model-First with tier filtering."""
        candidates: list[Candidate] = []

        for provider in self._config.providers:
            # Filter by allowed_providers
            if ctx.allowed_providers and provider.name not in ctx.allowed_providers:
                continue

            groups = self._get_groups(provider, ctx)

            for group in groups:
                for model in group.models:
                    for key in provider.keys:
                        candidates.append(
                            Candidate(
                                provider_name=provider.name,
                                model=model,
                                key_alias=key.alias,
                                key_token=key.token,
                                model_group=group.name,
                                base_url=provider.base_url,
                                client_type=provider.client_type,
                            )
                        )

        return candidates

    def _get_groups(
        self, provider: ProviderConfig, ctx: RoutingContext
    ) -> list[ModelGroupConfig]:
        """Get model groups filtered by RoutingContext."""
        if provider.model_groups:
            groups = list(provider.model_groups)
        elif provider.models:
            # Synthetic group for flat model lists
            groups = [
                ModelGroupConfig(name="_default", tier=1, models=provider.models)
            ]
        else:
            return []

        # Filter by model_group name
        if ctx.model_group:
            groups = [g for g in groups if g.name == ctx.model_group]
            if not ctx.allow_downgrade:
                return groups

        # Filter by tier ceiling
        groups = [g for g in groups if g.tier >= ctx.tier]

        # Sort by tier (best matching first)
        groups.sort(key=lambda g: g.tier)

        return groups

    def _get_provider_config(self, name: str) -> ProviderConfig | None:
        for p in self._config.providers:
            if p.name == name:
                return p
        return None

    async def _run_before_hooks(
        self, ctx: RoutingContext, candidate: Candidate
    ) -> bool:
        """Run before_request hooks. Returns False if any hook rejects."""
        for hook in self._hooks:
            if hasattr(hook, "before_request"):
                result = await hook.before_request(ctx, candidate)
                if result is False:
                    return False
        return True

    async def _run_after_hooks(
        self, ctx: RoutingContext, candidate: Candidate, usage: Usage
    ) -> None:
        """Run after_response hooks."""
        for hook in self._hooks:
            if hasattr(hook, "after_response"):
                await hook.after_response(ctx, candidate, usage)

    async def _run_fallback_hooks(
        self,
        ctx: RoutingContext,
        from_candidate: Candidate,
        all_candidates: list[Candidate],
        error: Exception,
    ) -> None:
        """Run on_fallback hooks."""
        for hook in self._hooks:
            if hasattr(hook, "on_fallback"):
                await hook.on_fallback(ctx, from_candidate, None, error)

    async def _emit_quota_warnings(self, warnings: list[QuotaWarning]) -> None:
        """Dispatch quota warnings to callback and hooks."""
        for warning in warnings:
            self._logger.warning(
                "Quota warning: %s at %.0f%% (%d/%d)",
                warning.scope,
                warning.percentage * 100,
                warning.current,
                warning.limit,
            )
            # Callback first
            if self._on_quota_warning is not None:
                await self._on_quota_warning(warning)
            # Then hooks
            for hook in self._hooks:
                if hasattr(hook, "on_quota_warning"):
                    await hook.on_quota_warning(warning)
