"""Invariant + seeded-simulation tests for the persistent-client pool (lever A of the TTFT work).
Proves the red-teamed safety properties hold by construction: fingerprint-gated reuse, pop-first
disposal, never-raising teardown, idle/LRU reclaim, and (seeded sim) that random op sequences never
reuse a stale client, never double-boot needlessly, and always recover a dead one. What the
fingerprint itself hashes lives in test_boot_fingerprint.py."""

import asyncio
import random
from typing import Dict, List

from backend.apps.agents.manager.run.client_pool import (
    ClientHandle,
    acquire_client,
    dispose_all_clients,
    dispose_client,
    dispose_client_soon,
    start_pool_sweeper,
    stop_pool_sweeper,
    trim_pool_to_cap,
)


class FakeClient:
    """Stands in for ClaudeSDKClient: counts connects/disconnects, can be killed, can raise on disconnect."""

    def __init__(self, registry: List["FakeClient"], raise_on_disconnect: bool = False):
        self.alive = True
        self.disconnected = False
        self.raise_on_disconnect = raise_on_disconnect
        registry.append(self)

    async def disconnect(self):
        self.disconnected = True
        self.alive = False
        if self.raise_on_disconnect:
            raise RuntimeError("teardown boom")


def test_reuse_respawn_force_and_teardown():
    async def run():
        pool: Dict[str, ClientHandle] = {}
        made: List[FakeClient] = []

        async def connect():
            return FakeClient(made)

        h1 = await acquire_client(pool, "s1", "fpA", connect)
        h2 = await acquire_client(pool, "s1", "fpA", connect)
        assert h1 is h2 and len(made) == 1

        h3 = await acquire_client(pool, "s1", "fpB", connect)
        assert h3 is not h1 and len(made) == 2 and made[0].disconnected

        h4 = await acquire_client(pool, "s1", "fpB", connect, force_respawn=True)
        assert h4 is not h3 and len(made) == 3 and made[1].disconnected

        await dispose_client(pool, "s1")
        assert "s1" not in pool and made[2].disconnected
        await dispose_client(pool, "s1")  # idempotent

        async def connect_bad():
            return FakeClient(made, raise_on_disconnect=True)

        await acquire_client(pool, "s2", "fp", connect_bad)
        await dispose_client(pool, "s2")  # teardown error swallowed
        assert "s2" not in pool

        await acquire_client(pool, "s3", "fp", connect)
        dispose_client_soon(pool, "s3")
        assert "s3" not in pool  # pop is sync-first
        await asyncio.sleep(0.01)
        assert made[-1].disconnected

        await acquire_client(pool, "s4", "fp", connect)
        await acquire_client(pool, "s5", "fp", connect)
        await dispose_all_clients(pool)
        assert not pool and all(c.disconnected for c in made)

    asyncio.run(run())


def test_idle_eviction():
    async def run():
        import backend.apps.agents.manager.run.client_pool as cp
        pool: Dict[str, ClientHandle] = {}
        made: List[FakeClient] = []

        async def connect():
            return FakeClient(made)

        old_ttl = cp.IDLE_EVICT_SECONDS
        cp.IDLE_EVICT_SECONDS = 0.05
        try:
            h = await acquire_client(pool, "s1", "fp", connect)
            await acquire_client(pool, "s2", "fp", connect)
            await asyncio.sleep(0.1)
            # s1 is mid-turn (lock held): the sweep must skip it and evict only the idle s2.
            async with h.lock:
                await cp.evict_idle_clients(pool)
                assert "s1" in pool and "s2" not in pool and made[1].disconnected
            await asyncio.sleep(0.1)
            await cp.evict_idle_clients(pool)
            assert "s1" not in pool and made[0].disconnected
            # a fresh acquire after eviction reconnects transparently
            h2 = await acquire_client(pool, "s1", "fp", connect)
            assert h2.client.alive
        finally:
            cp.IDLE_EVICT_SECONDS = old_ttl

    asyncio.run(run())


def test_cap_lru_eviction():
    """Over MAX_LIVE_CLIENTS, acquire trims the least-recently-used IDLE sessions and keeps the newest."""
    async def run():
        import backend.apps.agents.manager.run.client_pool as cp
        pool: Dict[str, ClientHandle] = {}
        made: List[FakeClient] = []

        async def connect():
            return FakeClient(made)

        old_max, old_guard = cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS
        cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = 3, 0.0
        try:
            for i in range(5):
                await acquire_client(pool, f"s{i}", "fp", connect)
                await asyncio.sleep(0.001)  # distinct last_used so LRU order is deterministic
            assert len(pool) == 3
            assert "s0" not in pool and "s1" not in pool  # two oldest reaped
            assert {"s2", "s3", "s4"} <= set(pool)
            assert made[0].disconnected and made[1].disconnected
            assert not made[3].disconnected and not made[4].disconnected
        finally:
            cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = old_max, old_guard

    asyncio.run(run())


def test_cap_soft_exceeds_when_busy():
    """A cap can't evict mid-turn clients: a new acquire over the cap exceeds it rather than kill a
    live turn, then trims back once they go idle."""
    async def run():
        import backend.apps.agents.manager.run.client_pool as cp
        pool: Dict[str, ClientHandle] = {}
        made: List[FakeClient] = []

        async def connect():
            return FakeClient(made)

        old_max, old_guard = cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS
        cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = 2, 0.4
        try:
            h0 = await acquire_client(pool, "s0", "fp", connect)
            h1 = await acquire_client(pool, "s1", "fp", connect)
            async with h0.lock, h1.lock:
                await acquire_client(pool, "s2", "fp", connect)
                # s0/s1 locked, s2 just-acquired (guard-protected): nothing is eligible, so the pool exceeds the cap.
                assert len(pool) == 3
                assert not made[0].disconnected and not made[1].disconnected
            await asyncio.sleep(0.5)  # past the guard: the now-idle sessions become eligible
            await trim_pool_to_cap(pool)
            assert len(pool) == 2 and made[0].disconnected  # oldest idle reaped back to cap
        finally:
            cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS = old_max, old_guard

    asyncio.run(run())


def test_pool_sweeper_reclaims_over_cap():
    """The background sweeper trims an over-cap pool on its timer, with no new turn to trigger it."""
    async def run():
        import backend.apps.agents.manager.run.client_pool as cp
        pool: Dict[str, ClientHandle] = {}
        made: List[FakeClient] = []

        async def connect():
            return FakeClient(made)

        old_max, old_guard, old_int = cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS, cp.SWEEP_INTERVAL_SECONDS
        cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS, cp.SWEEP_INTERVAL_SECONDS = 10, 0.0, 0.02
        try:
            for i in range(5):
                await acquire_client(pool, f"s{i}", "fp", connect)
                await asyncio.sleep(0.001)
            assert len(pool) == 5  # under the temporary high cap
            cp.MAX_LIVE_CLIENTS = 3
            task = start_pool_sweeper(pool)
            await asyncio.sleep(0.1)  # several sweep cycles
            await stop_pool_sweeper(task)
            assert len(pool) == 3
            assert "s0" not in pool and "s1" not in pool
            await stop_pool_sweeper(None)  # None is a no-op, must not raise
        finally:
            cp.MAX_LIVE_CLIENTS, cp.LRU_GUARD_SECONDS, cp.SWEEP_INTERVAL_SECONDS = old_max, old_guard, old_int

    asyncio.run(run())


def test_seeded_simulation_invariants():
    """Random op sequences: reuse only on identical fingerprint, dead clients always replaced, pool
    never re-serves a disposed client, and boots never exceed the one-shot baseline (one per turn)."""
    async def run():
        rng = random.Random(1337)
        pool: Dict[str, ClientHandle] = {}
        made: List[FakeClient] = []
        boots = 0
        turns = 0
        fp = "fp0"
        force = False

        async def connect():
            nonlocal boots
            boots += 1
            return FakeClient(made)

        for _ in range(300):
            op = rng.choice(["follow_up", "activate", "branch_or_fresh", "kill", "close"])
            if op == "follow_up":
                turns += 1
                h = await acquire_client(pool, "sim", fp, connect, force_respawn=force)
                force = False
                assert h.fingerprint == fp and not h.client.disconnected
                if not h.client.alive:  # dead client detected by the turn -> dispose + one respawn
                    await dispose_client(pool, "sim")
                    h = await acquire_client(pool, "sim", fp, connect)
                    assert h.client.alive
                async with h.lock:
                    assert h.lock.locked()  # single consumer while a turn drains
                    h.turns_served += 1
            elif op == "activate":
                fp = f"fp{rng.randint(0, 10**9)}"  # mcp_servers grew -> fingerprint changed
            elif op == "branch_or_fresh":
                force = True  # needs_fresh/fork read pre-build forces respawn
            elif op == "kill" and "sim" in pool:
                pool["sim"].client.alive = False
            elif op == "close":
                await dispose_client(pool, "sim")

        assert boots <= turns, f"persistent booted {boots}x for {turns} turns; one-shot baseline is {turns}"
        live = [c for c in made if not c.disconnected]
        assert len(live) <= 1, "at most the pooled client may be alive; everything else must be torn down"
        if "sim" in pool:
            assert not pool["sim"].client.disconnected

    asyncio.run(run())


def test_concurrent_acquires_share_one_spawn() -> None:
    """A pre-warm and a racing first turn must never double-spawn: the second spawn used to leak the first CLI."""
    import asyncio

    from backend.apps.agents.manager.run.client_pool import acquire_client

    async def run() -> None:
        pool: dict = {}
        spawns = 0

        class FakeClient:
            async def disconnect(self) -> None:
                return None

        async def connect_fn():
            nonlocal spawns
            spawns += 1
            await asyncio.sleep(0.05)
            return FakeClient()

        a, b = await asyncio.gather(
            acquire_client(pool, "sess-race", "fp1", connect_fn),
            acquire_client(pool, "sess-race", "fp1", connect_fn),
        )
        assert spawns == 1, f"double spawn: {spawns}"
        assert a is b
        assert pool["sess-race"].client is a.client

    asyncio.run(run())


def test_cancelled_waiter_does_not_kill_the_shared_spawn() -> None:
    import asyncio

    from backend.apps.agents.manager.run.client_pool import acquire_client

    async def run() -> None:
        pool: dict = {}

        class FakeClient:
            async def disconnect(self) -> None:
                return None

        async def connect_fn():
            await asyncio.sleep(0.08)
            return FakeClient()

        first = asyncio.ensure_future(acquire_client(pool, "sess-cancel", "fp1", connect_fn))
        await asyncio.sleep(0.01)
        second = asyncio.ensure_future(acquire_client(pool, "sess-cancel", "fp1", connect_fn))
        await asyncio.sleep(0.01)
        second.cancel()
        handle = await first
        assert pool["sess-cancel"].client is handle.client

    asyncio.run(run())