"""Reliability tests: SlotCap contention + fail-open, broker/subscription
resilience, SSE slot release on cancellation.

These drive the real code paths (no logic re-implemented in the test) and
pin the *documented* contracts read from the source:

  * ``SlotCap.acquire`` reserves via atomic INCR and rolls an over-cap
    reservation back with DECR; it fails **OPEN** (returns ``True``) on
    any backing-store exception (``incr`` / ``expire`` / ``decr``).
  * The broker / subscription layer has **no** built-in reconnect: a
    ``get_message`` raising propagates out of ``SSEStream`` (the
    ``finally`` still releases the slot + closes the sub). Recovery /
    retry is the consumer's supervisor's job — pinned here as a gap.
  * ``SSEStreamService`` releases the slot in a ``finally`` on the
    cancellation path too (client disconnect mid-stream).
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

import pytest

from mgf.livepush import SlotCap, SSEStreamService

if TYPE_CHECKING:
    from tests.conftest import FakeRedis


async def _never() -> bool:
    return False


# ---------------------------------------------------------------------------
# 1. SlotCap concurrency — no oversell under contention
# ---------------------------------------------------------------------------


class _YieldingRedis:
    """Async Redis double that faithfully models **atomic** INCR/DECR (real
    Redis is single-threaded: each command commits whole) while still
    yielding to the event loop, so concurrent ``acquire`` coroutines really
    interleave *between* SlotCap's separate awaits (incr → expire → decr).

    The read-modify-write is done in one synchronous step (no yield BETWEEN
    read and write) — that is the atomicity SlotCap's no-oversell design
    relies on. The yield is at the command boundary, where interleaving is
    legitimate. A check-then-act bug across SlotCap's own awaits would still
    surface here."""

    def __init__(self) -> None:
        self.store: dict[str, int] = {}

    async def incr(self, key: str) -> int:
        self.store[key] = self.store.get(key, 0) + 1  # atomic step
        result = self.store[key]
        await asyncio.sleep(0)  # yield at the command boundary
        return result

    async def decr(self, key: str) -> int:
        self.store[key] = self.store.get(key, 0) - 1  # atomic step
        result = self.store[key]
        await asyncio.sleep(0)
        return result

    async def expire(self, key: str, ttl: int) -> bool:
        await asyncio.sleep(0)
        return True

    async def set(self, key: str, value: int, ex: int | None = None) -> None:
        self.store[key] = value


@pytest.mark.concurrency
async def test_concurrent_acquire_never_oversells_one_key() -> None:
    cap_n = 5
    cap = SlotCap(_YieldingRedis(), cap=cap_n, key_prefix="t:")

    results = await asyncio.gather(*(cap.acquire("k") for _ in range(50)))

    granted = sum(results)
    assert granted <= cap_n, f"oversell: {granted} grants for cap={cap_n}"
    assert granted == cap_n, "every available slot should be handed out"


@pytest.mark.concurrency
async def test_release_under_contention_frees_slots_for_reuse() -> None:
    cap_n = 3
    redis = _YieldingRedis()
    cap = SlotCap(redis, cap=cap_n, key_prefix="t:")

    first = await asyncio.gather(*(cap.acquire("k") for _ in range(10)))
    assert sum(first) == cap_n
    assert redis.store["t:k"] == cap_n  # counter sits exactly at the cap

    # Release every granted slot concurrently, then re-acquire — the freed
    # slots must be reusable and the counter must floor cleanly at 0.
    await asyncio.gather(*(cap.release("k") for _ in range(cap_n)))
    assert redis.store["t:k"] == 0

    second = await asyncio.gather(*(cap.acquire("k") for _ in range(10)))
    assert sum(second) == cap_n


# ---------------------------------------------------------------------------
# 2. SlotCap fail-open semantics under backing-store error
# ---------------------------------------------------------------------------


class _BlowUpRedis:
    """Raises on whichever op is named in ``fail_on``; otherwise behaves."""

    def __init__(self, *, fail_on: set[str]) -> None:
        self.store: dict[str, int] = {}
        self.fail_on = fail_on

    async def incr(self, key: str) -> int:
        if "incr" in self.fail_on:
            raise RuntimeError("redis down: incr")
        self.store[key] = self.store.get(key, 0) + 1
        return self.store[key]

    async def decr(self, key: str) -> int:
        if "decr" in self.fail_on:
            raise RuntimeError("redis down: decr")
        self.store[key] = self.store.get(key, 0) - 1
        return self.store[key]

    async def expire(self, key: str, ttl: int) -> bool:
        if "expire" in self.fail_on:
            raise RuntimeError("redis down: expire")
        return True

    async def set(self, key: str, value: int, ex: int | None = None) -> None:
        if "set" in self.fail_on:
            raise RuntimeError("redis down: set")
        self.store[key] = value


@pytest.mark.parametrize("failing_op", ["incr", "expire", "decr"])
async def test_acquire_fails_open_on_any_backing_store_error(failing_op: str) -> None:
    """Contract (docstring): acquire returns True on ANY redis error — never
    raises, never denies. Covers the three ops acquire can touch:
    incr (always), expire (first acquire), decr (over-cap rollback)."""
    if failing_op == "decr":
        # decr only runs on the over-cap rollback path: prime to the cap with
        # a healthy store, then make decr blow up on the rollback attempt.
        redis = _BlowUpRedis(fail_on=set())
        cap = SlotCap(redis, cap=1, key_prefix="t:")
        assert await cap.acquire("k") is True  # fills the single slot
        redis.fail_on = {"decr"}
        # 2nd acquire is over-cap → tries to roll back via decr → that raises
        # → must fail OPEN rather than propagate.
        assert await cap.acquire("k") is True
    else:
        redis = _BlowUpRedis(fail_on={failing_op})
        cap = SlotCap(redis, cap=1, key_prefix="t:")
        assert await cap.acquire("k") is True


async def test_acquire_failopen_skips_on_acquire_callback() -> None:
    """On the fail-open path acquire returns early — observability callbacks
    do NOT fire (the slot was never truly reserved)."""
    fired: list[str] = []
    redis = _BlowUpRedis(fail_on={"incr"})
    cap = SlotCap(redis, cap=1, on_acquire=fired.append)
    assert await cap.acquire("k") is True
    assert fired == []  # callback skipped on the fail-open early return


async def test_release_swallows_backing_store_error() -> None:
    """release is best-effort: a redis error is swallowed, the on_release
    callback still fires (the slot is logically gone regardless)."""
    fired: list[str] = []
    redis = _BlowUpRedis(fail_on={"decr"})
    cap = SlotCap(redis, on_release=fired.append)
    await cap.release("k")  # must not raise despite decr blowing up
    assert fired == ["k"]


# ---------------------------------------------------------------------------
# 3. Broker reconnect / subscription resilience
# ---------------------------------------------------------------------------


class _FlakySubscription:
    """A subscription whose ``get_message`` raises ONCE then recovers — to
    pin whether anything in the stack retries. (It does not: the raise
    propagates; recovery is a consumer-supervisor concern.)"""

    def __init__(self, recovered_payload: bytes) -> None:
        self._raised = False
        self._payload = recovered_payload
        self.closed = False

    async def get_message(self, *, timeout: float) -> bytes | None:
        if not self._raised:
            self._raised = True
            raise ConnectionError("pubsub connection dropped")
        return self._payload

    async def close(self) -> None:
        self.closed = True


class _FlakyBroker:
    def __init__(self, sub: _FlakySubscription) -> None:
        self._sub = sub

    async def publish(self, channel: str, payload: str | bytes) -> int:
        return 1

    async def subscribe(self, channel: str) -> _FlakySubscription:
        return self._sub


async def test_subscription_error_propagates_and_slot_is_released(
    fake_redis: FakeRedis,
) -> None:
    """ACTUAL contract: there is NO reconnect in the broker/SSE layer. A
    ``get_message`` raising propagates out of the SSE iterator UNCHANGED
    (not swallowed, not retried). The ``finally`` still closes the sub and
    releases the slot — so the error is loud but leak-free. Recovery is the
    consumer's supervisor's job (documented gap, not a bug)."""
    sub = _FlakySubscription(recovered_payload=b"late")
    cap = SlotCap(fake_redis, key_prefix="t:")
    svc = SSEStreamService(_FlakyBroker(sub), cap, heartbeat_seconds=0.05)  # type: ignore[arg-type]

    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    gen = stream.body
    assert (await gen.__anext__()).startswith(b"retry:")  # slot acquired
    assert fake_redis.store["t:k"] == 1

    with pytest.raises(ConnectionError):
        await gen.__anext__()  # get_message raises; nothing retries it

    # The generator's finally ran: sub closed + slot released (no leak).
    assert sub.closed is True
    assert fake_redis.store["t:k"] == 0


async def test_redis_subscription_close_is_swallowing() -> None:
    """The Redis adapter's ``_RedisSubscription.close`` suppresses teardown
    errors so a broken connection can't crash stream cleanup. Pin that
    contract directly against the adapter."""
    from mgf.livepush._broker import _RedisSubscription

    class _BadPubSub:
        async def unsubscribe(self, channel: str) -> None:
            raise ConnectionError("already gone")

        async def aclose(self) -> None:
            raise ConnectionError("already gone")

    sub = _RedisSubscription(_BadPubSub(), "c")
    await sub.close()  # must not raise — errors are suppressed


# ---------------------------------------------------------------------------
# 4. SSE slot release on cancellation (client disconnect mid-stream)
# ---------------------------------------------------------------------------


class _BlockingSubscription:
    """``get_message`` blocks forever — models an idle stream a cancel must
    interrupt."""

    def __init__(self) -> None:
        self.closed = False

    async def get_message(self, *, timeout: float) -> bytes | None:
        await asyncio.Event().wait()  # never returns
        return None  # pragma: no cover

    async def close(self) -> None:
        self.closed = True


class _BlockingBroker:
    def __init__(self, sub: _BlockingSubscription) -> None:
        self._sub = sub

    async def publish(self, channel: str, payload: str | bytes) -> int:
        return 1

    async def subscribe(self, channel: str) -> _BlockingSubscription:
        return self._sub


async def test_slot_released_when_stream_cancelled_mid_stream(
    fake_redis: FakeRedis,
) -> None:
    """Cancellation path: a client disconnect cancels the body coroutine
    while it's awaiting the next message. The ``finally`` must close the sub
    and release the slot — no leak (distinct from the heartbeat/is_disconnected
    happy path already covered in test_sse.py)."""
    sub = _BlockingSubscription()
    cap = SlotCap(fake_redis, key_prefix="t:")
    svc = SSEStreamService(_BlockingBroker(sub), cap, heartbeat_seconds=30.0)  # type: ignore[arg-type]

    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    gen = stream.body
    assert (await gen.__anext__()).startswith(b"retry:")  # slot acquired
    assert fake_redis.store["t:k"] == 1

    # Drive the generator into its blocking get_message, then tear it down the
    # way an ASGI server does on client disconnect: cancel the in-flight pull,
    # which throws CancelledError into the suspended frame and runs the finally
    # (close sub + release). Then aclose() to fully retire the generator.
    pending = asyncio.ensure_future(gen.__anext__())
    await asyncio.sleep(0)  # let it reach the blocking await
    assert not pending.done()

    pending.cancel()
    with pytest.raises((asyncio.CancelledError, StopAsyncIteration)):
        await pending
    await gen.aclose()  # idempotent retire; generator already finished

    assert sub.closed is True
    assert fake_redis.store["t:k"] == 0  # slot freed on cancel — no leak


async def test_slot_released_on_explicit_cancellederror(
    fake_redis: FakeRedis,
) -> None:
    """ACTUAL contract: the source catches ``asyncio.CancelledError`` inside
    ``_event_gen`` and ``return``s, so the ``finally`` releases the slot and
    the cancellation is *absorbed* — the consuming ``async for`` ends cleanly
    rather than re-raising. Drive that exact branch by cancelling the task
    consuming the stream, then assert the slot was freed (no leak)."""
    sub = _BlockingSubscription()
    cap = SlotCap(fake_redis, key_prefix="t:")
    svc = SSEStreamService(_BlockingBroker(sub), cap, heartbeat_seconds=30.0)  # type: ignore[arg-type]

    stream = await svc.open_stream(channel="c", slot_key="k", is_disconnected=_never)
    gen = stream.body
    # Emit the retry frame first so the slot is observably acquired, then the
    # next pull blocks in get_message.
    assert (await gen.__anext__()).startswith(b"retry:")
    assert fake_redis.store["t:k"] == 1

    pull = asyncio.ensure_future(gen.__anext__())
    await asyncio.sleep(0)  # now blocked on get_message
    assert not pull.done()
    pull.cancel()  # CancelledError thrown into the suspended generator frame
    # The generator catches CancelledError and returns → __anext__ surfaces
    # StopAsyncIteration (the gen finished), not CancelledError.
    with pytest.raises((StopAsyncIteration, asyncio.CancelledError)):
        await pull

    assert sub.closed is True
    assert fake_redis.store["t:k"] == 0
