"""Unit tests for apex_stream.pool — health tracking + least-loaded selection."""

import time

import pytest

from apex_stream.pool import (
    ClientPool,
    NoClientAvailable,
)

CHUNK = b"x" * 1024


def _client(fail_times=0, name="c"):
    class FakeClient:
        def __init__(self):
            self.name = name
            self.fail_count = 0
            self.fail_times = fail_times

        async def fetch(self, n):
            if self.fail_count < self.fail_times:
                self.fail_count += 1
                raise ConnectionError("transport down")
            return CHUNK * n

    return FakeClient()


class TestPoolSelection:
    async def test_least_loaded_selected(self):
        pool = ClientPool(max_streams_per_client=4)
        c1 = _client(name="c1")
        c2 = _client(name="c2")
        await pool.register(1, c1)
        await pool.register(2, c2)
        h = await pool.select()
        assert h.client_id in (1, 2)
        h2 = await pool.select()
        assert h2.client_id != h.client_id or pool.max_streams_per_client == 1
        assert h.active_streams == h2.active_streams == 1

    async def test_capacity_exhausted_raises(self):
        pool = ClientPool(max_streams_per_client=1)
        await pool.register(1, _client())
        await pool.register(2, _client())
        await pool.select()
        await pool.select()
        with pytest.raises(NoClientAvailable):
            await pool.select()

    async def test_release_frees_slot(self):
        pool = ClientPool(max_streams_per_client=1)
        await pool.register(1, _client())
        await pool.register(2, _client())
        h = await pool.select()
        await pool.release(h.client_id)
        # now selectable again
        h2 = await pool.select()
        assert h2.client_id == h.client_id

    async def test_select_avoids_specific_client(self):
        pool = ClientPool(max_streams_per_client=4)
        await pool.register(1, _client(name="bad"))
        await pool.register(2, _client(name="good"))
        h = await pool.select(avoid=1)
        assert h.client_id == 2

    async def test_avoid_ignored_when_only_one_eligible(self):
        pool = ClientPool(max_streams_per_client=2)
        await pool.register(1, _client(name="only"))
        h = await pool.select(avoid=1)
        assert h.client_id == 1

    async def test_least_loaded_prefers_idle(self):
        pool = ClientPool(max_streams_per_client=10)
        await pool.register(1, _client(name="busy"))
        await pool.register(2, _client(name="idle"))
        await pool.select()  # loads client 1 (or 2)
        idle = pool.select()
        # whichever was least loaded should be chosen again or the idle one
        h2 = await idle
        snapshot = pool.snapshot()["clients"]
        loads = {c["id"]: c["active_streams"] for c in snapshot}
        assert max(loads.values()) <= 2

    async def test_note_success_resets_failures(self):
        pool = ClientPool(max_streams_per_client=1)
        await pool.register(1, _client())
        await pool.note_failure(1, telegram_error=True)
        await pool.note_failure(1, telegram_error=True)
        await pool.note_success(1)
        assert pool.snapshot()["clients"][0]["consecutive_failures"] == 0

    async def test_cancellation_is_neutral(self):
        pool = ClientPool(max_streams_per_client=1, failure_threshold=2)
        await pool.register(1, _client())
        await pool.note_failure(1, telegram_error=False)  # cancellation
        assert pool.snapshot()["clients"][0]["state"] == "healthy"
        assert pool.snapshot()["clients"][0]["consecutive_failures"] == 0

    async def test_rate_limit_uses_longer_cooldown(self):
        pool = ClientPool(
            max_streams_per_client=1,
            failure_threshold=1,
            cooldown_seconds=5,
        )
        await pool.register(1, _client())
        await pool.note_failure(1, telegram_error=True, cooldown_seconds=30)
        client = pool.snapshot()["clients"][0]
        assert client["cooldown_remaining"] > 5
        assert client["rate_limit_count"] == 1


class TestPoolDcPreference:
    async def test_prefers_matching_dc_client(self):
        pool = ClientPool(max_streams_per_client=4)
        await pool.register(1, _client(name="dc2"), dc_id=2)
        await pool.register(2, _client(name="dc5"), dc_id=5)
        h = await pool.select(dc_id=5)
        assert h.client_id == 2

    async def test_falls_back_when_no_dc_match(self):
        pool = ClientPool(max_streams_per_client=2)
        await pool.register(1, _client(name="dc2"), dc_id=2)
        h = await pool.select(dc_id=4)
        assert h.client_id == 1

    async def test_matching_dc_beats_lower_load(self):
        pool = ClientPool(max_streams_per_client=8)
        await pool.register(2, _client(name="dc5-lite"), dc_id=5)
        await pool.register(1, _client(name="dc2-busy"), dc_id=2)
        h = await pool.select(dc_id=5)
        assert h.client_id == 2

    async def test_avoid_still_respected_among_matching(self):
        pool = ClientPool(max_streams_per_client=4)
        await pool.register(2, _client(name="dc5-b"), dc_id=5)
        await pool.register(3, _client(name="dc5-c"), dc_id=5)
        await pool.register(1, _client(name="dc2"), dc_id=2)
        h = await pool.select(dc_id=5, avoid=2)
        assert h.client_id == 3

    async def test_dc_field_in_snapshot_and_register_roundtrip(self):
        pool = ClientPool(max_streams_per_client=1)
        await pool.register(7, _client(), dc_id=3)
        h = await pool.select()
        assert h.dc_id == 3
        assert pool.snapshot()["clients"][0]["dc_id"] == 3


class TestPoolCooldown:
    async def test_cooldown_after_threshold(self):
        pool = ClientPool(failure_threshold=3, cooldown_seconds=60)
        await pool.register(1, _client())
        for _ in range(3):
            await pool.note_failure(1, telegram_error=True)
        snap = pool.snapshot()["clients"][0]
        assert snap["state"] == "cooldown"
        assert snap["cooldown_remaining"] > 0

    async def test_eligible_after_cooldown(self, monkeypatch):
        pool = ClientPool(failure_threshold=1, cooldown_seconds=1)
        await pool.register(1, _client())
        await pool.note_failure(1, telegram_error=True)
        assert pool.snapshot()["clients"][0]["state"] == "cooldown"
        # re-open after cooldown
        handle = pool._clients[1]
        handle.cooldown_until = time.monotonic() - 0.1
        h = await pool.select()
        assert h.client_id == 1
        state = pool.snapshot()["clients"][0]["state"]
        assert state != "unavailable"


class TestPoolSnapshots:
    async def test_snapshot_exposes_workload_and_health_counters(self):
        pool = ClientPool(max_streams_per_client=2, failure_threshold=2)
        await pool.register(1, _client())
        await pool.select()
        await pool.note_failure(1, telegram_error=True)
        await pool.note_failure(1, telegram_error=True)
        await pool.note_success(1)

        snap = pool.snapshot()
        assert snap["active_streams"] == 1
        assert snap["counters"] == {"successes": 1, "failures": 2, "cooldowns": 1}
        client = snap["clients"][0]
        assert client["success_count"] == 1
        assert client["failure_count"] == 2
        assert client["cooldown_count"] == 1

    async def test_unregister_removes_client(self):
        pool = ClientPool()
        await pool.register(1, _client())
        await pool.register(2, _client())
        assert len(pool) == 2
        await pool.unregister(1)
        assert len(pool) == 1