from __future__ import annotations

from typing import TYPE_CHECKING

import pytest

from mgf.livepush import CapExceededError, MockBroker, SlotCap, SSEStreamService

if TYPE_CHECKING:
    from tests.conftest import FakeRedis


async def _never() -> bool:
    return False


async def test_first_frame_is_retry_directive(fake_redis: FakeRedis) -> None:
    svc = SSEStreamService(
        MockBroker(), SlotCap(fake_redis, key_prefix="t:"), heartbeat_seconds=0.05
    )
    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    gen = stream.body
    first = await gen.__anext__()
    assert first.startswith(b"retry:")
    assert stream.headers["Cache-Control"] == "no-cache, no-transform"
    assert stream.headers["X-Accel-Buffering"] == "no"
    await gen.aclose()


async def test_data_frame_delivered(fake_redis: FakeRedis) -> None:
    broker = MockBroker()
    svc = SSEStreamService(broker, SlotCap(fake_redis, key_prefix="t:"), heartbeat_seconds=0.5)
    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    gen = stream.body
    assert (await gen.__anext__()).startswith(b"retry:")
    await broker.publish("c", "ping")
    assert await gen.__anext__() == b"data: ping\n\n"
    await gen.aclose()


async def test_keepalive_on_heartbeat_timeout(fake_redis: FakeRedis) -> None:
    svc = SSEStreamService(
        MockBroker(), SlotCap(fake_redis, key_prefix="t:"), heartbeat_seconds=0.02
    )
    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    gen = stream.body
    await gen.__anext__()  # retry
    assert await gen.__anext__() == b": keep-alive\n\n"
    await gen.aclose()


async def test_cap_exceeded_raises_before_subscribe(fake_redis: FakeRedis) -> None:
    svc = SSEStreamService(
        MockBroker(), SlotCap(fake_redis, cap=1, key_prefix="t:"), heartbeat_seconds=0.05
    )
    await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    with pytest.raises(CapExceededError):
        await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)


async def test_slot_released_on_disconnect(fake_redis: FakeRedis) -> None:
    cap = SlotCap(fake_redis, key_prefix="t:")
    svc = SSEStreamService(MockBroker(), cap, heartbeat_seconds=0.02)
    disconnected = False

    async def is_disc() -> bool:
        return disconnected

    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=is_disc)
    gen = stream.body
    await gen.__anext__()  # retry — slot acquired
    assert fake_redis.store["t:k"] == 1
    disconnected = True
    with pytest.raises(StopAsyncIteration):
        await gen.__anext__()  # heartbeat timeout → disconnected → finally releases
    assert fake_redis.store["t:k"] == 0
