"""Provider capacity manager: simulated capacity only. No production provider quota."""

from __future__ import annotations

import asyncio
import inspect
import time
from types import SimpleNamespace

import pytest
from src.core.provider_capacity import (
    MemoryCapacityStore,
    ProviderCapacityManager,
    ProviderCapacityUnavailableError,
    RedisCapacityStore,
    capacity_bucket,
    override_capacity_store,
    policy_from_settings,
)
from src.core.provider_capacity.algorithm import available_units, try_reserve
from src.core.provider_capacity.types import (
    BucketSnapshot,
    CapacityObservation,
    CapacityPolicy,
)
from src.integrations.openai_capacity import (
    estimate_openai_request_tokens,
    observations_from_openai_headers,
    parse_openai_reset_duration,
)
from src.workers.tasks.project_discovery import reserve_discovery_call

POLICY = CapacityPolicy(
    safety_fraction=0.9,
    wait_seconds=1.0,
    unknown_max_in_flight=2,
    stale_observation_seconds=120,
    poll_seconds=0.01,
)


def bucket(name="tokens"):
    return capacity_bucket(provider="openai", secret="test-key", model="sim", dimension=name)


def obs(now_ms: int, *, limit: int, remaining: int, reset: float = 60) -> CapacityObservation:
    return CapacityObservation(
        "tokens",
        limit=limit,
        remaining=remaining,
        reset_in_seconds=reset,
        observed_at_ms=now_ms,
    )


def observed(state: BucketSnapshot, limit: int, remaining: int, *, reset_in=60, now_ms=1_000_000):
    state.observed_limit = limit
    state.observed_remaining = remaining
    state.observed_at_ms = now_ms
    state.reset_at_ms = now_ms + int(reset_in * 1000)
    return state


@pytest.mark.asyncio
async def test_capacity_below_safety_threshold_executes_immediately():
    store = MemoryCapacityStore()
    manager = ProviderCapacityManager(store, POLICY)
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        CapacityObservation(
            "tokens", limit=100_000, remaining=90_000, reset_in_seconds=60, observed_at_ms=now
        ),
        now,
    )
    trace: dict = {}
    reservation = await manager.acquire(bucket(), 1_000, trace=trace)
    assert reservation.decision.granted is True
    assert reservation.wait_ms < 50
    assert trace["provider_capacity"]["capacity_granted"] is True
    assert trace["provider_capacity"]["capacity_queued"] is False
    await manager.release(reservation)


@pytest.mark.asyncio
async def test_capacity_above_safety_threshold_is_queued():
    store = MemoryCapacityStore()
    manager = ProviderCapacityManager(store, POLICY)
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        obs(now, limit=10_000, remaining=10_000),
        now,
    )
    first = await manager.acquire(bucket(), 8_000)
    assert first.decision.granted is True
    queued = asyncio.create_task(manager.acquire(bucket(), 3_000, trace={}))
    await asyncio.sleep(0.05)
    assert queued.done() is False
    await manager.release(first)
    second = await queued
    assert second.wait_ms >= 40
    assert second.decision.granted is True
    await manager.release(second)


@pytest.mark.asyncio
async def test_capacity_reset_lets_queued_work_proceed():
    store = MemoryCapacityStore()
    manager = ProviderCapacityManager(store, POLICY)
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        CapacityObservation(
            "tokens",
            limit=20_000,
            remaining=0,
            reset_in_seconds=-1,
            observed_at_ms=now,
        ),
        now,
    )
    reservation = await manager.acquire(bucket(), 1_000)
    assert reservation.decision.mode == "reset_refilled"
    assert reservation.decision.granted is True
    await manager.release(reservation)


@pytest.mark.asyncio
async def test_concurrent_workers_cannot_over_reserve_bucket():
    store = MemoryCapacityStore()
    manager = ProviderCapacityManager(store, POLICY)
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        obs(now, limit=10_000, remaining=10_000),
        now,
    )

    async def take(demand: int):
        reservation = await manager.acquire(bucket(), demand)
        await asyncio.sleep(0.05)
        await manager.release(reservation)
        return reservation

    first, second = await asyncio.gather(take(6_000), take(6_000), return_exceptions=True)
    granted = [item for item in (first, second) if not isinstance(item, Exception)]
    # Safe capacity is 9000; two 6000 reservations cannot both be in-flight.
    assert len(granted) == 2
    overlapping = first.wait_ms == 0 and second.wait_ms == 0
    assert overlapping is False


@pytest.mark.asyncio
async def test_queue_deadline_is_bounded_capacity_failure():
    store = MemoryCapacityStore()
    policy = CapacityPolicy(**{**POLICY.__dict__, "wait_seconds": 0.05})
    manager = ProviderCapacityManager(store, policy)
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        obs(now, limit=10_000, remaining=100),
        now,
    )
    with pytest.raises(ProviderCapacityUnavailableError):
        await manager.acquire(bucket(), 5_000)


@pytest.mark.asyncio
async def test_provider_feedback_reduces_available_capacity():
    store = MemoryCapacityStore()
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        obs(now, limit=10_000, remaining=8_000),
        now,
    )
    await store.reconcile(
        bucket(),
        obs(now, limit=10_000, remaining=500),
        now,
    )
    available, mode, safe = available_units(
        BucketSnapshot.from_dict(
            {
                "observed_limit": 10_000,
                "observed_remaining": 500,
                "observed_at_ms": now,
                "reset_at_ms": now + 60_000,
            }
        ),
        POLICY,
        now,
    )
    assert mode == "observed"
    assert safe == 9_000
    assert available == 500
    manager = ProviderCapacityManager(
        store, CapacityPolicy(**{**POLICY.__dict__, "wait_seconds": 0.05})
    )
    with pytest.raises(ProviderCapacityUnavailableError):
        await manager.acquire(bucket(), 2_000)


def test_provider_remaining_constrains_local_reservations():
    """Org-scoped remaining can be below local headroom; that remaining wins."""
    now = 1_000_000
    state = observed(BucketSnapshot(), 200_000, 3_031, now_ms=now)
    available, mode, safe = available_units(state, POLICY, now)
    assert mode == "observed"
    assert safe == 180_000
    assert available == 3_031
    denied = try_reserve(state, demand=6_569, waiter_id="a", now_ms=now, policy=POLICY)
    assert denied.granted is False


def test_missing_headers_use_unknown_in_flight_fallback():
    state = BucketSnapshot()
    available, mode, safe = available_units(state, POLICY, 1_000)
    assert mode == "unknown_fallback"
    assert safe is None
    assert available > 0
    state.in_flight = 2
    available, mode, _ = available_units(state, POLICY, 1_000)
    assert available == 0


def test_algorithm_fifo_does_not_let_later_work_jump_the_queue():
    state = observed(BucketSnapshot(), 10_000, 10_000)
    first = try_reserve(state, demand=8_000, waiter_id="a", now_ms=1_000_000, policy=POLICY)
    assert first.granted is True
    second = try_reserve(state, demand=2_000, waiter_id="b", now_ms=1_000_001, policy=POLICY)
    assert second.granted is False
    third = try_reserve(state, demand=2_000, waiter_id="c", now_ms=1_000_002, policy=POLICY)
    assert third.granted is False
    assert [item[0] for item in state.waiters] == ["b", "c"]


def test_openai_headers_are_parsed_without_fabricating_missing_fields():
    observations = observations_from_openai_headers(
        {
            "x-ratelimit-limit-tokens": "200000",
            "x-ratelimit-remaining-tokens": "3031",
            "x-ratelimit-reset-tokens": "59.09s",
            "x-ratelimit-limit-requests": "10000",
            "x-ratelimit-remaining-requests": "9736",
            "x-ratelimit-reset-requests": "37m59.817s",
        }
    )
    by_dim = {item.dimension: item for item in observations}
    assert by_dim["tokens"].limit == 200000
    assert by_dim["tokens"].remaining == 3031
    assert by_dim["tokens"].reset_in_seconds == pytest.approx(59.09)
    assert by_dim["requests"].limit == 10000
    assert observations_from_openai_headers({}) == []
    assert parse_openai_reset_duration("6m0s") == 360
    assert estimate_openai_request_tokens(
        [{"role": "user", "content": "hello"}], max_completion_tokens=10
    ) > 10


def test_sixty_per_minute_policy_is_unchanged_and_independent():
    source = inspect.getsource(reserve_discovery_call)
    assert "project-intake:budget:" in source
    assert "project_intake_calls_per_minute" in source
    assert "provider_capacity" not in source
    from src.config import Settings

    assert Settings.model_fields["project_intake_calls_per_minute"].default == 60
    assert Settings.model_fields["project_intake_timeout_seconds"].default == 8
    assert Settings.model_fields["project_discovery_timeout_seconds"].default == 25
    assert Settings.model_fields["provider_capacity_safety_fraction"].default == 0.90


def test_provider_timeout_constant_unchanged():
    from src.integrations import project_intake_client as client

    source = inspect.getsource(client._project_structured_call)
    assert "timeout=timeout" in source
    assert "max_retries=0" in source
    assert "provider_capacity" in source


@pytest.mark.asyncio
async def test_redis_store_mutex_prevents_double_grant(monkeypatch):
    class FakeRedis:
        def __init__(self):
            self.data: dict = {}

        async def get(self, key):
            return self.data.get(key)

        async def set(self, key, value, nx=False, px=None, ex=None):
            if nx and key in self.data:
                return False
            self.data[key] = value
            return True

        async def eval(self, script, n, *args):
            key, token = args[0], args[1]
            if self.data.get(key) == token:
                self.data.pop(key, None)
                return 1
            return 0

        async def __aenter__(self):
            return self

        async def __aexit__(self, *exc):
            return False

    shared = FakeRedis()
    store = RedisCapacityStore("redis://test", redis_factory=lambda: shared)
    now = int(time.time() * 1000)
    await store.reconcile(
        bucket(),
        obs(now, limit=10_000, remaining=10_000),
        now,
    )
    first, _ = await store.try_reserve(bucket(), 6_000, "w1", now, POLICY)
    second, snapshot = await store.try_reserve(bucket(), 6_000, "w2", now, POLICY)
    assert first.granted is True
    assert second.granted is False
    assert snapshot.reserved == 6_000


def test_policy_comes_from_settings_not_hardcoded_tpm():
    settings = SimpleNamespace(
        provider_capacity_safety_fraction=0.5,
        provider_capacity_wait_seconds=12,
        provider_capacity_unknown_max_in_flight=3,
        provider_capacity_stale_observation_seconds=30,
    )
    policy = policy_from_settings(settings)
    assert policy.safety_fraction == 0.5
    assert policy.wait_seconds == 12
    source = inspect.getsource(available_units)
    assert "200000" not in source
    assert "180000" not in source


@pytest.mark.asyncio
async def test_override_store_is_used_by_manager_factory():
    store = MemoryCapacityStore()
    override_capacity_store(store)
    from src.core.provider_capacity import manager_from_settings

    manager = manager_from_settings(SimpleNamespace(redis_url="redis://ignored"))
    assert manager.store is store