"""Tests for lifecycle hooks."""

from __future__ import annotations

from typing import Any

import pytest

from llm_rotator._types import Candidate, LLMResponse, RoutingContext, Usage
from llm_rotator.config import RotatorConfig
from llm_rotator.exceptions import AllAttemptsFailedError, ModelRateLimitError
from llm_rotator.rotator import LLMRotator
from tests.conftest import FakeLLMClient


def _resp(model: str = "gpt-4o", content: str = "ok", total: int = 20) -> LLMResponse:
    return LLMResponse(
        content=content,
        usage=Usage(
            prompt_tokens=total // 2,
            completion_tokens=total - total // 2,
            total_tokens=total,
        ),
        model=model,
        provider="fake",
        key_alias="unknown",
    )


def _simple_config() -> RotatorConfig:
    return RotatorConfig(
        providers=[
            {
                "name": "openai",
                "client_type": "openai",
                "priority": 1,
                "models": ["gpt-4o", "gpt-5"],
                "keys": [
                    {"token": "sk-1", "alias": "key1"},
                    {"token": "sk-2", "alias": "key2"},
                ],
            }
        ]
    )


class RecordingHook:
    """Records all hook calls for assertion."""

    def __init__(
        self,
        *,
        reject_candidates: set[str] | None = None,
        raise_on_fallback: Exception | None = None,
    ) -> None:
        self.before_calls: list[tuple[RoutingContext, Candidate]] = []
        self.after_calls: list[tuple[RoutingContext, Candidate, Usage]] = []
        self.fallback_calls: list[dict[str, Any]] = []
        self._reject = reject_candidates or set()
        self._raise_on_fallback = raise_on_fallback

    async def before_request(
        self, context: RoutingContext, candidate: Candidate
    ) -> bool:
        self.before_calls.append((context, candidate))
        return candidate.key_alias not in self._reject

    async def after_response(
        self, context: RoutingContext, candidate: Candidate, usage: Usage
    ) -> None:
        self.after_calls.append((context, candidate, usage))

    async def on_fallback(
        self,
        context: RoutingContext,
        from_candidate: Candidate,
        to_candidate: Candidate | None,
        error: Exception,
    ) -> None:
        self.fallback_calls.append({
            "context": context,
            "from": from_candidate,
            "to": to_candidate,
            "error": error,
        })
        if self._raise_on_fallback:
            raise self._raise_on_fallback


class TestBeforeRequestHook:
    async def test_skip_candidate(self):
        """Hook returns False → candidate skipped, next one used."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(_resp(content="from key2"))

        hook = RecordingHook(reject_candidates={"key1"})
        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook)

        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}]
        )

        assert result.content == "from key2"
        assert client.calls[0]["api_key"] == "sk-2"
        assert len(hook.before_calls) == 2  # called for key1 (rejected) and key2

    async def test_all_rejected(self):
        """Hook rejects all candidates → AllAttemptsFailedError."""
        config = _simple_config()
        client = FakeLLMClient()

        hook = RecordingHook(reject_candidates={"key1", "key2"})
        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook)

        with pytest.raises(AllAttemptsFailedError):
            await rotator.complete(
                messages=[{"role": "user", "content": "hi"}]
            )

        assert len(client.calls) == 0  # no HTTP calls made


class TestAfterResponseHook:
    async def test_receives_usage(self):
        """after_response gets correct usage data."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(_resp(total=150))

        hook = RecordingHook()
        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook)

        await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert len(hook.after_calls) == 1
        _, _, usage = hook.after_calls[0]
        assert usage.total_tokens == 150


class TestOnFallbackHook:
    async def test_on_fallback_called(self):
        """on_fallback is called when switching candidates."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(
            ModelRateLimitError(key_alias="key1", model="gpt-4o", retry_after=60)
        )
        client.enqueue(_resp(content="key2 ok"))

        hook = RecordingHook()
        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook)

        await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert len(hook.fallback_calls) >= 1
        fb = hook.fallback_calls[0]
        assert fb["from"].key_alias == "key1"
        assert isinstance(fb["error"], ModelRateLimitError)

    async def test_on_fallback_exception_aborts(self):
        """If on_fallback raises, rotation chain aborts."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(
            ModelRateLimitError(key_alias="key1", model="gpt-4o", retry_after=60)
        )
        client.enqueue(_resp(content="should not reach"))

        hook = RecordingHook(raise_on_fallback=RuntimeError("abort!"))
        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook)

        with pytest.raises(RuntimeError, match="abort!"):
            await rotator.complete(
                messages=[{"role": "user", "content": "hi"}]
            )


class TestHooksExecutionOrder:
    async def test_order_before_then_after(self):
        """before_request → HTTP → after_response."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(_resp())

        order: list[str] = []

        class OrderHook:
            async def before_request(self, ctx: RoutingContext, c: Candidate) -> bool:
                order.append("before")
                return True

            async def after_response(
                self, ctx: RoutingContext, c: Candidate, u: Usage
            ) -> None:
                order.append("after")

        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(OrderHook())

        await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert order == ["before", "after"]


class TestMultipleHooks:
    async def test_multiple_hooks_all_called(self):
        """All registered hooks are called."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(_resp())

        hook1 = RecordingHook()
        hook2 = RecordingHook()

        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook1)
        rotator.add_hook(hook2)

        await rotator.complete(messages=[{"role": "user", "content": "hi"}])

        assert len(hook1.after_calls) == 1
        assert len(hook2.after_calls) == 1

    async def test_any_hook_rejects_skips_candidate(self):
        """If any hook rejects, candidate is skipped."""
        config = _simple_config()
        client = FakeLLMClient()
        client.enqueue(_resp(content="from key2"))

        hook_allow = RecordingHook()  # allows all
        hook_reject = RecordingHook(reject_candidates={"key1"})

        rotator = LLMRotator(config, clients={"openai": client})
        rotator.add_hook(hook_allow)
        rotator.add_hook(hook_reject)

        result = await rotator.complete(
            messages=[{"role": "user", "content": "hi"}]
        )
        assert result.content == "from key2"
