"""Tests for Issue #3787: Probe-worker pool tail-latency.

Focus on:
- Per-home worker pool size N=2
- Non-blocking acquire returns an idle worker when available
- Returns None immediately when all workers are locked (no blocking wait)
- Scoped invalidation only flushes the active home's workers
- Full invalidation (no provider_id) flushes all workers
- Cleanup iterates nested worker lists correctly
- Lock held across selector handoff prevents TOCTOU races
"""

import threading
import time
import unittest
from pathlib import Path
from unittest import mock

from api.providers import (
    _ACCOUNT_USAGE_WORKERS_PER_HOME,
    _account_usage_worker_pool,
    _account_usage_worker_pool_lock,
    _cleanup_account_usage_probe_workers,
    _close_account_usage_probe_workers,
    _get_account_usage_probe_worker,
    invalidate_account_usage_status_cache,
)


class TestProbeWorkerPoolPerHome(unittest.TestCase):
    """Test per-home worker pool with N=2 configuration."""

    def setUp(self):
        """Clear the pool before each test."""
        with _account_usage_worker_pool_lock:
            _account_usage_worker_pool.clear()

    def tearDown(self):
        """Clean up after each test."""
        with _account_usage_worker_pool_lock:
            _account_usage_worker_pool.clear()

    def test_pool_creates_two_workers_per_home_key(self):
        """Pool should create exactly N=2 workers for each home key."""
        home = Path.home() / ".hermes"
        worker = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(worker)
        worker._lock.release()

        with _account_usage_worker_pool_lock:
            key = str(Path(home))
            workers_list = _account_usage_worker_pool.get(key)
            self.assertIsNotNone(workers_list)
            self.assertEqual(len(workers_list), _ACCOUNT_USAGE_WORKERS_PER_HOME)
            self.assertEqual(_ACCOUNT_USAGE_WORKERS_PER_HOME, 2)

    def test_nonblocking_acquire_returns_idle_worker(self):
        """Non-blocking acquire should return an idle worker when one is free."""
        home = Path("/tmp/test_nonblocking_idle")

        # Populate the pool
        worker_initial = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(worker_initial)
        worker_initial._lock.release()

        with _account_usage_worker_pool_lock:
            key = str(Path(home))
            workers = _account_usage_worker_pool[key]
            self.assertEqual(len(workers), 2)

        # Hold workers[0] locked from a background thread to simulate actual use
        lock_holder = threading.Event()
        release_signal = threading.Event()

        def hold_lock():
            workers[0]._lock.acquire()
            lock_holder.set()
            release_signal.wait(timeout=2.0)
            workers[0]._lock.release()

        thread = threading.Thread(target=hold_lock, daemon=True)
        thread.start()
        lock_holder.wait(timeout=2.0)

        try:
            worker = _get_account_usage_probe_worker(home)
            self.assertIsNotNone(worker)
            self.assertIs(worker, workers[1])
            worker._lock.release()
        finally:
            release_signal.set()
            thread.join(timeout=1.0)

    def test_returns_none_when_all_workers_locked(self):
        """Should return None immediately when all workers are locked from other threads."""
        home = Path("/tmp/test_immediate_none")

        # Populate the pool
        worker_initial = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(worker_initial)
        worker_initial._lock.release()

        with _account_usage_worker_pool_lock:
            key = str(Path(home))
            workers = _account_usage_worker_pool[key]

        # Hold both workers locked from a background thread
        lock_holder = threading.Event()
        release_signal = threading.Event()

        def hold_locks():
            workers[0]._lock.acquire()
            workers[1]._lock.acquire()
            lock_holder.set()
            release_signal.wait(timeout=2.0)
            workers[1]._lock.release()
            workers[0]._lock.release()

        thread = threading.Thread(target=hold_locks, daemon=True)
        thread.start()
        lock_holder.wait(timeout=2.0)

        try:
            start = time.time()
            result = _get_account_usage_probe_worker(home)
            elapsed = time.time() - start

            self.assertIsNone(result)
            # Should return immediately, not block
            self.assertLess(elapsed, 0.1)
        finally:
            release_signal.set()
            thread.join(timeout=1.0)

    def test_scoped_invalidation_only_flushes_active_home(self):
        """Scoped invalidation with provider_id should only flush active home."""
        home1 = Path.home() / ".hermes"
        home2 = Path("/tmp/other_home")

        worker1 = _get_account_usage_probe_worker(home1)
        worker2 = _get_account_usage_probe_worker(home2)

        self.assertIsNotNone(worker1)
        self.assertIsNotNone(worker2)
        worker1._lock.release()
        worker2._lock.release()

        with _account_usage_worker_pool_lock:
            initial_count = len(_account_usage_worker_pool)
            self.assertEqual(initial_count, 2)

        with mock.patch("api.providers._get_hermes_home", return_value=home1):
            invalidate_account_usage_status_cache(provider_id="anthropic")
            time.sleep(0.2)

        with _account_usage_worker_pool_lock:
            remaining_keys = list(_account_usage_worker_pool.keys())
            self.assertNotIn(str(Path(home1)), remaining_keys)
            self.assertIn(str(Path(home2)), remaining_keys)

    def test_full_invalidation_flushes_all_workers(self):
        """Full invalidation (no provider_id) should flush all workers."""
        home1 = Path.home() / ".hermes"
        home2 = Path("/tmp/other_home")

        worker1 = _get_account_usage_probe_worker(home1)
        worker2 = _get_account_usage_probe_worker(home2)

        self.assertIsNotNone(worker1)
        self.assertIsNotNone(worker2)
        worker1._lock.release()
        worker2._lock.release()

        with _account_usage_worker_pool_lock:
            self.assertEqual(len(_account_usage_worker_pool), 2)

        invalidate_account_usage_status_cache(provider_id=None)
        time.sleep(0.2)

        with _account_usage_worker_pool_lock:
            self.assertEqual(len(_account_usage_worker_pool), 0)

    def test_cleanup_iterates_nested_worker_lists(self):
        """Cleanup should properly iterate over nested worker lists."""
        home = Path.home() / ".hermes"

        worker = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(worker)
        worker._lock.release()

        with _account_usage_worker_pool_lock:
            key = str(Path(home))
            workers_list = _account_usage_worker_pool.get(key)
            self.assertEqual(len(workers_list), 2)

        now = time.monotonic() + 100000
        _cleanup_account_usage_probe_workers(now=now, idle_seconds=1.0)

        with _account_usage_worker_pool_lock:
            if _account_usage_worker_pool.get(str(Path(home))):
                remaining = _account_usage_worker_pool.get(str(Path(home)), [])
                self.assertEqual(len(remaining), 0)
            else:
                self.assertNotIn(str(Path(home)), _account_usage_worker_pool)

    def test_partial_cleanup_replenishes_pool(self):
        """When cleanup removes one stale worker but the other is busy, pool replenishes to N=2."""
        home = Path("/tmp/test_replenish")

        worker = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(worker)
        worker._lock.release()

        with _account_usage_worker_pool_lock:
            key = str(Path(home))
            workers = _account_usage_worker_pool[key]
            self.assertEqual(len(workers), 2)

        # Hold workers[1] locked from a background thread (simulating active use)
        lock_holder = threading.Event()
        release_signal = threading.Event()

        def hold_lock():
            workers[1]._lock.acquire()
            lock_holder.set()
            release_signal.wait(timeout=5.0)
            workers[1]._lock.release()

        thread = threading.Thread(target=hold_lock, daemon=True)
        thread.start()
        lock_holder.wait(timeout=2.0)

        try:
            # Run cleanup far in the future; workers[0] is idle and stale,
            # workers[1] is locked so it survives
            now = time.monotonic() + 100000
            _cleanup_account_usage_probe_workers(now=now, idle_seconds=1.0)

            with _account_usage_worker_pool_lock:
                remaining = _account_usage_worker_pool.get(key, [])
                # Pool should be replenished back to 2
                self.assertEqual(len(remaining), _ACCOUNT_USAGE_WORKERS_PER_HOME)
                # The original busy worker should still be present
                self.assertIn(workers[1], remaining)
        finally:
            release_signal.set()
            thread.join(timeout=1.0)

    def test_synchronous_close_flattens_nested_lists(self):
        """Synchronous close should flatten nested lists correctly."""
        home = Path.home() / ".hermes"

        worker = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(worker)
        worker._lock.release()

        with _account_usage_worker_pool_lock:
            key = str(Path(home))
            workers_list = _account_usage_worker_pool.get(key)
            self.assertEqual(len(workers_list), 2)

        _close_account_usage_probe_workers()

        with _account_usage_worker_pool_lock:
            self.assertEqual(len(_account_usage_worker_pool), 0)

    def test_concurrent_selector_no_double_claim(self):
        """Two threads holding returned workers simultaneously never share the same instance."""
        home = Path("/tmp/test_concurrent")

        # Pre-populate pool
        initial = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(initial)
        initial._lock.release()

        with _account_usage_worker_pool_lock:
            workers = _account_usage_worker_pool[str(Path(home))]

        # Lock workers[0] from a background thread so only workers[1] is free
        lock_holder = threading.Event()
        release_signal = threading.Event()

        def hold_first():
            workers[0]._lock.acquire()
            lock_holder.set()
            release_signal.wait(timeout=5.0)
            workers[0]._lock.release()

        holder_thread = threading.Thread(target=hold_first, daemon=True)
        holder_thread.start()
        lock_holder.wait(timeout=2.0)

        # Two threads race; each holds its worker until both have finished grabbing
        results = [None, None]
        barrier = threading.Barrier(2, timeout=2.0)
        both_done = threading.Barrier(2, timeout=2.0)

        def grab(idx):
            barrier.wait()
            w = _get_account_usage_probe_worker(home)
            results[idx] = w
            try:
                both_done.wait()
            except threading.BrokenBarrierError:
                pass
            if w is not None:
                w._lock.release()

        t0 = threading.Thread(target=grab, args=(0,), daemon=True)
        t1 = threading.Thread(target=grab, args=(1,), daemon=True)
        t0.start()
        t1.start()
        t0.join(timeout=3.0)
        t1.join(timeout=3.0)

        try:
            got = [r for r in results if r is not None]
            # At most one thread should have gotten a worker (workers[1]);
            # the other gets None because the lock is already held
            self.assertEqual(len(got), 1, "Expected exactly one winner, got %d" % len(got))
            self.assertIs(got[0], workers[1])
        finally:
            release_signal.set()
            holder_thread.join(timeout=1.0)

    def test_invalidation_cannot_pop_worker_between_lookup_and_claim(self):
        """Getter keeps the pool lock through worker claim, so invalidation cannot pop first."""
        home = Path("/tmp/test_claim_under_pool_lock")

        initial = _get_account_usage_probe_worker(home)
        self.assertIsNotNone(initial)
        initial._lock.release()

        with _account_usage_worker_pool_lock:
            workers = _account_usage_worker_pool[str(Path(home))]

        workers[1]._lock.acquire()
        self.addCleanup(workers[1]._lock.release)

        acquire_entered = threading.Event()
        allow_acquire = threading.Event()
        claimed = threading.Event()
        release_claim = threading.Event()
        invalidation_done = threading.Event()
        result: dict[str, object] = {}
        target = workers[0]
        original_lock = target._lock
        self_outer = self

        class GateLock:
            def __init__(self, inner):
                self._inner = inner

            def acquire(self, blocking=True, timeout=-1):
                if blocking is False:
                    self_outer.assertTrue(
                        _account_usage_worker_pool_lock.locked(),
                        "worker claim must happen while the pool lock is still held",
                    )
                    acquire_entered.set()
                    allow_acquire.wait(timeout=2.0)
                if timeout == -1:
                    return self._inner.acquire(blocking)
                return self._inner.acquire(blocking, timeout)

            def release(self):
                return self._inner.release()

            def __enter__(self):
                self._inner.acquire()
                return self

            def __exit__(self, exc_type, exc, tb):
                self._inner.release()
                return False

        target._lock = GateLock(original_lock)

        def run_getter():
            result["worker"] = _get_account_usage_probe_worker(home)
            claimed.set()
            release_claim.wait(timeout=2.0)
            if result["worker"] is not None:
                result["worker"]._lock.release()

        def run_invalidation():
            with mock.patch("api.providers._get_hermes_home", return_value=home):
                invalidate_account_usage_status_cache(provider_id="anthropic")
            invalidation_done.set()

        fetch_thread = threading.Thread(target=run_getter, daemon=True)
        invalidation_thread = threading.Thread(target=run_invalidation, daemon=True)
        fetch_thread.start()
        self.assertTrue(acquire_entered.wait(timeout=2.0))

        invalidation_thread.start()
        time.sleep(0.05)
        self.assertFalse(invalidation_done.is_set())

        allow_acquire.set()
        self.assertTrue(claimed.wait(timeout=2.0))
        self.assertIs(result["worker"], target)

        release_claim.set()
        fetch_thread.join(timeout=2.0)

        invalidation_thread.join(timeout=2.0)
        self.assertTrue(invalidation_done.is_set())
        with _account_usage_worker_pool_lock:
            self.assertNotIn(str(Path(home)), _account_usage_worker_pool)


class TestWorkerPoolConfiguration(unittest.TestCase):
    """Test that constants are properly configured."""

    def test_workers_per_home_is_two(self):
        """Should have exactly 2 workers per home."""
        self.assertEqual(_ACCOUNT_USAGE_WORKERS_PER_HOME, 2)


if __name__ == "__main__":
    unittest.main()