"""Unit tests for apex_stream.stream — end-to-end range streaming."""

import asyncio

import pytest

from apex_stream.cache import BoundedChunkCache
from apex_stream.config import Config
from apex_stream.errors import SourceError
from apex_stream.inflight import InFlightRegistry
from apex_stream.metrics import StreamMetrics
from apex_stream.pool import ClientPool
from apex_stream.ranges import RangeNotSatisfiable
from apex_stream.stream import FileInfo, StreamDriver


class _FakeTelegram:
    """Simulates one Pyrogram client exposing stream_media()-like reads."""

    def __init__(self, file_size, fail_call_count=0, chunk_size=1024 * 1024):
        self.file_size = file_size
        self.fail_call_count = fail_call_count  # fail first N read_range calls
        self.chunk_size = chunk_size
        self.fetch_count = 0

    async def read_range(self, offset_chunk, chunk_count):
        if self.fetch_count < self.fail_call_count:
            self.fetch_count += 1
            raise ConnectionError("telegram transport down")
        for i in range(chunk_count):
            self.fetch_count += 1
            gidx = offset_chunk + i
            if gidx >= (self.file_size + self.chunk_size - 1) // self.chunk_size:
                return
            # deterministic bytes: identify chunk by global index
            b = bytearray(self.chunk_size)
            for j in range(0, self.chunk_size, 64):
                tag = f"c{gidx:06d}".encode()
                b[j : j + len(tag)] = tag
            yield (i, bytes(b))


def _make_driver(
    file_size=10 * 1024 * 1024,
    failures=0,
    cache_mb=32,
    run_size=8,
    max_retries=1,
    num_clients=1,
    max_global=32,
    max_per_source=8,
):
    cfg = Config(
        {
            "APEX_STREAM_CHUNK_SIZE": "1048576",
            "APEX_STREAM_RUN_CHUNKS": str(run_size),
            "APEX_STREAM_RUN_WINDOW": "4",
            "APEX_STREAM_MAX_GLOBAL": str(max_global),
            "APEX_STREAM_MAX_PER_CLIENT": str(num_clients * 4),
            "APEX_STREAM_MAX_PER_SOURCE": str(max_per_source),
            "APEX_STREAM_MAX_INFLIGHT": "16",
            "APEX_STREAM_CACHE_ENABLED": "true",
            "APEX_STREAM_CACHE_MB": str(cache_mb),
            "APEX_STREAM_CACHE_TTL_SECONDS": "60",
            "APEX_STREAM_EXTRA_TOKENS": "",
            "APEX_TELEGRAM_CLIENT_COOLDOWN": "5",
            "APEX_TELEGRAM_FAILURE_THRESHOLD": "2",
            "APEX_TELEGRAM_MAX_RETRIES": str(max_retries),
            "APEX_TELEGRAM_MAX_CONCURRENT": "4",
        }
    )
    pool = ClientPool(
        max_streams_per_client=cfg.max_per_client,
        failure_threshold=cfg.client_failure_threshold,
        cooldown_seconds=cfg.client_cooldown_seconds,
    )
    chunk_size = cfg.chunk_size
    total_chunks = (file_size + chunk_size - 1) // chunk_size

    async def resolve_fn(chat_id, message_id):
        return FileInfo(size=file_size, name=None, mime_type="video/mp4")

    async def fetch_fn(client, chat_id, message_id, run_start, chunk_count):
        # client is _FakeTelegram
        async for off, chunk in client.read_range(run_start, chunk_count):
            yield (off, chunk)

    cache = BoundedChunkCache(
        max_bytes=cfg.cache_size_bytes if cache_mb else 0,
        ttl_seconds=cfg.cache_ttl_seconds,
    )
    registry = InFlightRegistry(
        cache=cache, run_size=cfg.run_size, max_window=cfg.run_window
    )
    metrics = StreamMetrics()
    driver = StreamDriver(cfg, pool, registry, metrics, resolve_fn, fetch_fn)
    return driver, pool, registry, cache, metrics, file_size


def _expected_bytes(global_chunk, chunk_size):
    b = bytearray(chunk_size)
    for j in range(0, chunk_size, 64):
        tag = f"c{global_chunk:06d}".encode()
        b[j : j + len(tag)] = tag
    return bytes(b)


class TestPlan:
    async def test_no_range_full_file(self):
        driver, *_ = _make_driver()
        plan = await driver.plan(1, 2, None)
        assert plan.status == 200
        assert plan.content_length == plan.file_size
        assert plan.content_range is None
        assert plan.layout.total_needed == plan.file_size

    async def test_partial_range_206(self):
        driver, *_ = _make_driver()
        plan = await driver.plan(1, 2, "bytes=0-1048575")
        assert plan.status == 206
        assert plan.content_length == 1048576
        assert plan.content_range == f"bytes 0-1048575/{10 * 1024 * 1024}"

    async def test_range_past_eof_416(self):
        driver, *_ra = _make_driver()
        with pytest.raises(RangeNotSatisfiable):
            await driver.plan(1, 2, "bytes=99999999-")


class TestGenerate:
    async def test_full_file_stream(self):
        driver, pool, *_ = _make_driver(file_size=3 * 1024 * 1024)
        await pool.register(1, _FakeTelegram(file_size=3 * 1024 * 1024))
        plan = await driver.plan(1, 2, None)
        chunks = [c async for c in driver.generate(plan)]
        body = b"".join(chunks)
        assert len(body) == plan.file_size == 3 * 1024 * 1024

    async def test_range_stream_trims_edges(self):
        driver, pool, *_ = _make_driver()
        await pool.register(1, _FakeTelegram(file_size=10 * 1024 * 1024))
        # seek into the middle of chunk 0 and stop mid-chunk 2
        cs = 1024 * 1024
        start, end = 5000, 2 * cs + 4000
        plan = await driver.plan(1, 2, f"bytes={start}-{end}")
        body = b"".join([c async for c in driver.generate(plan)])
        assert len(body) == end - start + 1
        # first bytes belong to chunk 0 at offset 5000
        expected_first = _expected_bytes(0, cs)[5000:5000 + 64]
        assert body[:64] == expected_first
        # trailing bytes end exactly at `end` inside chunk 2
        within_chunk2 = end - 2 * cs
        expected_last = _expected_bytes(2, cs)[within_chunk2 - 63 : within_chunk2 + 1]
        assert body[-64:] == expected_last

    async def test_two_clients_stream_different_ranges(self):
        driver, pool, *_ = _make_driver(num_clients=2)
        await pool.register(1, _FakeTelegram(file_size=10 * 1024 * 1024))
        await pool.register(2, _FakeTelegram(file_size=10 * 1024 * 1024))
        p1 = await driver.plan(1, 2, "bytes=0-1048575")
        p2 = await driver.plan(1, 2, "bytes=1048576-2097151")
        b1, b2 = await asyncio.gather(
            _collect(driver, p1), _collect(driver, p2)
        )
        assert len(b1) == 1048576
        assert len(b2) == 1048576
        assert b1 == _expected_bytes(0, 1024 * 1024)
        assert b2 == _expected_bytes(1, 1024 * 1024)


async def _collect(driver, plan):
    return b"".join([c async for c in driver.generate(plan)])


class TestFailover:
    async def test_retry_after_transport_failure(self):
        driver, pool, registry, *_ = _make_driver(
            file_size=8 * 1024 * 1024, max_retries=2, num_clients=1
        )
        fake = _FakeTelegram(file_size=8 * 1024 * 1024, fail_call_count=1)
        await pool.register(1, fake)
        plan = await driver.plan(1, 2, None)
        chunks = [c async for c in driver.generate(plan)]
        body = b"".join(chunks)
        assert len(body) == 8 * 1024 * 1024  # retry recovered the full file
        assert pool.snapshot()["clients"][0]["consecutive_failures"] == 0

    async def test_retries_exhausted_raise(self):
        driver, pool, *_ = _make_driver(
            file_size=8 * 1024 * 1024, max_retries=1, num_clients=1
        )
        fake = _FakeTelegram(file_size=8 * 1024 * 1024, fail_call_count=999)
        await pool.register(1, fake)
        plan = await driver.plan(1, 2, None)
        with pytest.raises(ConnectionError):
            _ = [c async for c in driver.generate(plan)]

    async def test_cancellation_releases_client(self):
        driver, pool, *_ = _make_driver(num_clients=1)
        await pool.register(1, _FakeTelegram(file_size=4 * 1024 * 1024))
        plan = await driver.plan(1, 2, None)

        async def consume_until_cancel():
            gen = driver.generate(plan)
            async for chunk in gen:
                break  # simulate early FFmpeg stop
            await gen.aclose()

        await consume_until_cancel()
        # the client's stream counters should be back to 0
        assert pool.snapshot()["clients"][0]["active_streams"] == 0
        assert pool.snapshot()["clients"][0]["active_requests"] == 0


class _FakeSourceBroken:
    """read_range() that always raises a source-level error (e.g. expired ref)."""

    def __init__(self):
        self.chunk_size = 1024 * 1024

    async def read_range(self, offset_chunk, chunk_count):
        for _ in ():  # make this an async generator that fails on first iteration
            yield None
        raise SourceError("media reference expired and refresh failed")


class _FakeShortFile:
    """Serves only half of each requested run (file shorter than recorded)."""

    def __init__(self, chunks_served=4):
        self.chunks_served = chunks_served
        self.chunk_size = 1024 * 1024

    async def read_range(self, offset_chunk, chunk_count):
        for i in range(min(self.chunks_served, chunk_count)):
            b = bytearray(self.chunk_size)
            yield (i, bytes(b))


class TestSourceErrorClassification:
    """Source-level failures must fail over / surface WITHOUT cooling down the client."""

    async def test_source_error_does_not_cooldown(self):
        driver, pool, *_ = _make_driver(
            file_size=8 * 1024 * 1024, max_retries=2, num_clients=1
        )
        await pool.register(1, _FakeSourceBroken())
        plan = await driver.plan(1, 2, None)
        with pytest.raises(SourceError):
            _ = [c async for c in driver.generate(plan)]
        snap = pool.snapshot()["clients"][0]
        assert snap["state"] == "healthy"
        assert snap["consecutive_failures"] == 0

    async def test_source_error_fails_over_to_healthy_client(self):
        driver, pool, *_ = _make_driver(
            file_size=2 * 1024 * 1024, max_retries=1, num_clients=2
        )
        await pool.register(1, _FakeSourceBroken())
        await pool.register(2, _FakeTelegram(file_size=2 * 1024 * 1024))
        plan = await driver.plan(1, 2, None)
        body = b"".join([c async for c in driver.generate(plan)])
        assert len(body) == 2 * 1024 * 1024  # healthy client 2 served full
        states = {c["id"]: c["state"] for c in pool.snapshot()["clients"]}
        assert states[1] == "healthy"
        assert states[2] == "healthy"

    async def test_short_file_raises_stream_exhausted_without_cooldown(self):
        driver, pool, *_ = _make_driver(
            file_size=8 * 1024 * 1024, max_retries=2, num_clients=1
        )
        await pool.register(1, _FakeShortFile(chunks_served=4))
        plan = await driver.plan(1, 2, None)
        with pytest.raises(Exception):  # StreamExhausted surfaces to caller
            _ = [c async for c in driver.generate(plan)]
        snap = pool.snapshot()["clients"][0]
        assert snap["state"] == "healthy"
        assert snap["consecutive_failures"] == 0

    async def test_transport_failure_still_cools_down(self):
        driver, pool, *_ = _make_driver(
            file_size=8 * 1024 * 1024,
            max_retries=1,
            num_clients=1,
        )
        await pool.register(1, _FakeTelegram(file_size=8 * 1024 * 1024, fail_call_count=999))
        plan = await driver.plan(1, 2, None)
        with pytest.raises(ConnectionError):
            _ = [c async for c in driver.generate(plan)]
        snap = pool.snapshot()["clients"][0]
        # a real transport failure is still recorded (cooldown candidate)
        assert snap["consecutive_failures"] == 1
        assert snap["state"] == "degraded"


class TestAdmission:
    """Global + per-source concurrency caps are actually enforced."""

    async def test_global_limit_blocks_third_stream(self):
        driver, pool, *_ = _make_driver(
            file_size=1024 * 1024, max_global=2, max_per_source=8, num_clients=1
        )
        await pool.register(1, _FakeTelegram(file_size=1024 * 1024))

        async def _first_chunk(gen):
            async for _chunk in gen:
                return

        g1 = driver.generate(await driver.plan(1, 1, None))
        g2 = driver.generate(await driver.plan(2, 1, None))
        g3 = driver.generate(await driver.plan(3, 1, None))

        await asyncio.wait_for(_first_chunk(g1), 1.0)
        await asyncio.wait_for(_first_chunk(g2), 1.0)

        # a third concurrent stream must stay pending on the global gate.
        # Spawn it as a background task and verify it has not started the
        # generator body before a short deadline (without cancelling it).
        t3 = asyncio.ensure_future(_first_chunk(g3))
        done, _ = await asyncio.wait({t3}, timeout=0.2)
        assert t3 not in done, "third stream ran despite max_global_streams=2"
        assert driver.status_snapshot()["admission"]["active_streams"] == 2
        t3.cancel()
        try:
            await t3
        except asyncio.CancelledError:
            pass
        await g1.aclose()
        await g2.aclose()

    async def test_global_limit_releases_on_cancel(self):
        driver, pool, *_ = _make_driver(
            file_size=1024 * 1024, max_global=1, max_per_source=8, num_clients=1
        )
        await pool.register(1, _FakeTelegram(file_size=1024 * 1024))

        async def _first_chunk(gen):
            async for _chunk in gen:
                return

        g1 = driver.generate(await driver.plan(1, 1, None))
        g2 = driver.generate(await driver.plan(2, 1, None))

        await asyncio.wait_for(_first_chunk(g1), 1.0)
        t2 = asyncio.ensure_future(_first_chunk(g2))
        done, _ = await asyncio.wait({t2}, timeout=0.2)
        assert t2 not in done, "second stream ran despite max_global_streams=1"

        # cancelling g1 (e.g. ffmpeg stop) frees the global slot for g2
        await g1.aclose()
        await asyncio.wait_for(t2, 1.0)
        assert driver.status_snapshot()["admission"]["active_streams"] == 1
        await g2.aclose()

    async def test_per_source_limit_independent_of_other_sources(self):
        driver, pool, *_ = _make_driver(
            file_size=1024 * 1024, max_global=8, max_per_source=1, num_clients=1
        )
        await pool.register(1, _FakeTelegram(file_size=1024 * 1024))

        async def _first_chunk(gen):
            async for _chunk in gen:
                return

        # two different sources stream concurrently (per-source cap is 1, but
        # global cap permits both)
        ga = driver.generate(await driver.plan(10, 1, None))
        gb = driver.generate(await driver.plan(11, 1, None))
        await asyncio.wait_for(_first_chunk(ga), 1.0)
        await asyncio.wait_for(_first_chunk(gb), 1.0)
        assert driver.status_snapshot()["admission"]["active_streams"] == 2

        # a second stream on the same source must wait for the first
        gb2 = driver.generate(await driver.plan(11, 1, None))
        tb2 = asyncio.ensure_future(_first_chunk(gb2))
        done, _ = await asyncio.wait({tb2}, timeout=0.2)
        assert tb2 not in done, "second source stream ran despite max_per_source=1"

        await gb.aclose()
        await asyncio.wait_for(tb2, 1.0)
        await ga.aclose()
        await gb2.aclose()