"""Tests for quota warning notifications (Phase 14)."""

from __future__ import annotations

import pytest

from llm_rotator._types import Candidate, Usage
from llm_rotator.backends import InMemoryBackend
from llm_rotator.config import (
    ClientType,
    KeyConfig,
    ModelGroupConfig,
    ProviderConfig,
    RequestQuotaConfig,
    ResetSchedule,
    RotatorConfig,
    TokenQuotaConfig,
)
from llm_rotator.quota import QuotaManager, QuotaWarning

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_candidate(
    provider: str = "openai",
    model: str = "gpt-4o",
    key_alias: str = "main_key",
    group: str = "flagship",
) -> Candidate:
    return Candidate(
        provider_name=provider,
        model=model,
        key_alias=key_alias,
        key_token="sk-test",
        model_group=group,
        base_url=None,
        client_type=ClientType.OPENAI,
    )


def _make_provider(
    name: str = "openai",
    group_name: str = "flagship",
    token_limit: int = 100_000,
    key_alias: str = "main_key",
    request_quota: RequestQuotaConfig | None = None,
) -> ProviderConfig:
    return ProviderConfig(
        name=name,
        client_type=ClientType.OPENAI,
        priority=1,
        model_groups=[
            ModelGroupConfig(
                name=group_name,
                tier=1,
                models=["gpt-4o"],
                token_quota=TokenQuotaConfig(
                    limit=token_limit,
                    reset=ResetSchedule.DAILY_UTC,
                ),
            ),
        ],
        keys=[
            KeyConfig(token="sk-test", alias=key_alias, request_quota=request_quota),
        ],
    )


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_warning_at_80_percent() -> None:
    """At 80% usage, record_usage returns a warning."""
    backend = InMemoryBackend()
    provider = _make_provider(token_limit=100_000)
    qm = QuotaManager(backend, [provider], warning_threshold=0.8)

    candidate = _make_candidate()

    # Pre-load 79_000 tokens (79%)
    await backend.increment_quota("openai/flagship", 79_000, 86400)

    # Record 1000 more → 80_000 (80%) — crosses threshold
    usage = Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000)
    warnings = await qm.record_usage(candidate, usage)

    assert len(warnings) == 1
    w = warnings[0]
    assert w.scope == "openai/flagship"
    assert w.current == 80_000
    assert w.limit == 100_000
    assert abs(w.percentage - 0.80) < 0.01


@pytest.mark.asyncio
async def test_no_warning_below_threshold() -> None:
    """At 79% usage, no warning is emitted."""
    backend = InMemoryBackend()
    provider = _make_provider(token_limit=100_000)
    qm = QuotaManager(backend, [provider], warning_threshold=0.8)

    candidate = _make_candidate()

    # Pre-load 78_000 tokens
    await backend.increment_quota("openai/flagship", 78_000, 86400)

    # Record 1000 more → 79_000 (79%)
    usage = Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000)
    warnings = await qm.record_usage(candidate, usage)

    assert warnings == []


@pytest.mark.asyncio
async def test_warning_custom_threshold() -> None:
    """Custom threshold=0.9 triggers at 90%."""
    backend = InMemoryBackend()
    provider = _make_provider(token_limit=100_000)
    qm = QuotaManager(backend, [provider], warning_threshold=0.9)

    candidate = _make_candidate()
    await backend.increment_quota("openai/flagship", 89_000, 86400)

    # Record 1000 → 90_000 (90%)
    usage = Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000)
    warnings = await qm.record_usage(candidate, usage)

    assert len(warnings) == 1
    assert abs(warnings[0].percentage - 0.90) < 0.01


@pytest.mark.asyncio
async def test_warning_called_once_per_cross() -> None:
    """Warning fires once per threshold crossing, not on every request."""
    backend = InMemoryBackend()
    provider = _make_provider(token_limit=100_000)
    qm = QuotaManager(backend, [provider], warning_threshold=0.8)

    candidate = _make_candidate()
    await backend.increment_quota("openai/flagship", 79_000, 86400)

    usage = Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000)

    # First request crosses threshold → warning
    warnings1 = await qm.record_usage(candidate, usage)
    assert len(warnings1) == 1

    # Second request still above threshold → no warning
    warnings2 = await qm.record_usage(candidate, usage)
    assert warnings2 == []


@pytest.mark.asyncio
async def test_warning_receives_quota_info() -> None:
    """Warning contains scope, current, limit, percentage."""
    backend = InMemoryBackend()
    provider = _make_provider(token_limit=250_000)
    qm = QuotaManager(backend, [provider], warning_threshold=0.8)

    candidate = _make_candidate()
    await backend.increment_quota("openai/flagship", 199_000, 86400)

    usage = Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000)
    warnings = await qm.record_usage(candidate, usage)

    assert len(warnings) == 1
    w = warnings[0]
    assert w.scope == "openai/flagship"
    assert w.current == 200_000
    assert w.limit == 250_000
    assert abs(w.percentage - 0.80) < 0.01


@pytest.mark.asyncio
async def test_warning_for_token_and_request_quotas() -> None:
    """Warning works for both token and request quotas."""
    backend = InMemoryBackend()
    provider = _make_provider(
        token_limit=100_000,
        request_quota=RequestQuotaConfig(
            limit=100,
            reset=ResetSchedule.DAILY_UTC,
        ),
    )
    qm = QuotaManager(backend, [provider], warning_threshold=0.8)

    candidate = _make_candidate()

    # Pre-load token quota to 79% and request quota to 79%
    await backend.increment_quota("openai/flagship", 79_000, 86400)
    await backend.increment_quota("openai/main_key/requests", 79, 86400)

    # Record usage — crosses threshold for both
    usage = Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000)
    warnings = await qm.record_usage(candidate, usage)

    assert len(warnings) == 2
    scopes = {w.scope for w in warnings}
    assert "openai/flagship" in scopes
    assert "openai/main_key/requests" in scopes


@pytest.mark.asyncio
async def test_no_callback_no_error() -> None:
    """Without a registered callback, warnings are returned but no error."""
    backend = InMemoryBackend()
    provider = _make_provider(token_limit=100)
    qm = QuotaManager(backend, [provider], warning_threshold=0.8)

    candidate = _make_candidate()
    await backend.increment_quota("openai/flagship", 79, 86400)

    usage = Usage(prompt_tokens=10, completion_tokens=11, total_tokens=21)
    # Should not raise
    warnings = await qm.record_usage(candidate, usage)
    assert len(warnings) == 1


@pytest.mark.asyncio
async def test_simple_callback_in_constructor() -> None:
    """on_quota_warning callback in LLMRotator constructor is called."""
    from unittest.mock import AsyncMock

    from llm_rotator.rotator import LLMRotator

    received: list[QuotaWarning] = []

    async def on_warning(warning: QuotaWarning) -> None:
        received.append(warning)

    provider = _make_provider(token_limit=100_000)
    config = RotatorConfig(providers=[provider])
    backend = InMemoryBackend()
    await backend.increment_quota("openai/flagship", 79_000, 86400)

    rotator = LLMRotator(
        config,
        backend=backend,
        on_quota_warning=on_warning,
    )

    # Mock client
    mock_client = AsyncMock()
    mock_client.generate.return_value = _make_response()
    rotator._clients["openai"] = mock_client

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

    assert len(received) == 1
    assert received[0].scope == "openai/flagship"


@pytest.mark.asyncio
async def test_hook_on_quota_warning() -> None:
    """Lifecycle hook on_quota_warning method is called."""
    from unittest.mock import AsyncMock

    from llm_rotator.rotator import LLMRotator

    provider = _make_provider(token_limit=100_000)
    config = RotatorConfig(providers=[provider])
    backend = InMemoryBackend()
    await backend.increment_quota("openai/flagship", 79_000, 86400)

    rotator = LLMRotator(config, backend=backend)

    hook = AsyncMock()
    hook.before_request = AsyncMock(return_value=True)
    hook.after_response = AsyncMock()
    hook.on_quota_warning = AsyncMock()
    rotator.add_hook(hook)

    mock_client = AsyncMock()
    mock_client.generate.return_value = _make_response()
    rotator._clients["openai"] = mock_client

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

    hook.on_quota_warning.assert_called_once()
    warning = hook.on_quota_warning.call_args[0][0]
    assert warning.scope == "openai/flagship"


@pytest.mark.asyncio
async def test_both_callback_and_hook() -> None:
    """Both callback and hook are called; callback first."""
    from unittest.mock import AsyncMock

    from llm_rotator.rotator import LLMRotator

    call_order: list[str] = []

    async def on_warning(warning: QuotaWarning) -> None:
        call_order.append("callback")

    provider = _make_provider(token_limit=100_000)
    config = RotatorConfig(providers=[provider])
    backend = InMemoryBackend()
    await backend.increment_quota("openai/flagship", 79_000, 86400)

    rotator = LLMRotator(config, backend=backend, on_quota_warning=on_warning)

    hook = AsyncMock()
    hook.before_request = AsyncMock(return_value=True)
    hook.after_response = AsyncMock()

    async def hook_on_warning(warning: QuotaWarning) -> None:
        call_order.append("hook")

    hook.on_quota_warning = hook_on_warning
    rotator.add_hook(hook)

    mock_client = AsyncMock()
    mock_client.generate.return_value = _make_response()
    rotator._clients["openai"] = mock_client

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

    assert call_order == ["callback", "hook"]


# ---------------------------------------------------------------------------
# Helpers for rotator-level tests
# ---------------------------------------------------------------------------


def _make_response():
    from llm_rotator._types import LLMResponse

    return LLMResponse(
        content="Hello!",
        usage=Usage(prompt_tokens=500, completion_tokens=500, total_tokens=1000),
        model="gpt-4o",
        provider="openai",
        key_alias="main_key",
    )
