import threading
import time
import unittest

from token_bucket import TokenBucket


class TestTokenBucket(unittest.TestCase):
    def test_basic_acquire_and_refill(self):
        b = TokenBucket(rate=10.0, capacity=5.0)
        self.assertTrue(b.try_acquire(5.0))
        self.assertFalse(b.try_acquire(1.0))
        time.sleep(0.15)
        self.assertTrue(b.try_acquire(1.0))

    def test_capacity_cap(self):
        b = TokenBucket(rate=100.0, capacity=3.0)
        time.sleep(0.1)
        self.assertLessEqual(b.tokens(), 3.0 + 1e-9)

    def test_acquire_blocks_and_returns(self):
        b = TokenBucket(rate=20.0, capacity=1.0)
        t0 = time.monotonic()
        b.try_acquire(1.0)
        ok = b.acquire(1.0, timeout=0.2)
        dt = time.monotonic() - t0
        self.assertTrue(ok)
        self.assertGreaterEqual(dt, 0.04)

    def test_n_exceeds_capacity_raises(self):
        b = TokenBucket(rate=1.0, capacity=1.0)
        with self.assertRaises(ValueError):
            b.acquire(2.0, timeout=0.01)
        with self.assertRaises(ValueError):
            b.try_acquire(2.0)

    def test_invalid_params(self):
        for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
            with self.assertRaises(ValueError):
                TokenBucket(**kwargs)
        with self.assertRaises(ValueError):
            TokenBucket(rate=1, capacity=1).try_acquire(0)

    def test_concurrent_no_overconsume(self):
        b = TokenBucket(rate=1e-9, capacity=10.0)
        successes = {"n": 0}
        lock = threading.Lock()

        def worker():
            if b.try_acquire(1.0):
                with lock:
                    successes["n"] += 1

        threads = [threading.Thread(target=worker) for _ in range(50)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        self.assertEqual(successes["n"], 10)

    def test_acquire_timeout_false(self):
        b = TokenBucket(rate=1.0, capacity=1.0)
        b.try_acquire(1.0)
        t0 = time.monotonic()
        ok = b.acquire(1.0, timeout=0.05)
        dt = time.monotonic() - t0
        self.assertFalse(ok)
        self.assertLess(dt, 0.2)

    def test_reset(self):
        b = TokenBucket(rate=1.0, capacity=2.0)
        b.try_acquire(2.0)
        self.assertFalse(b.try_acquire(1.0))
        b.reset()
        self.assertTrue(b.try_acquire(2.0))

    def test_partial_never_consumes(self):
        b = TokenBucket(rate=1.0, capacity=2.0)
        b.try_acquire(2.0)
        self.assertFalse(b.try_acquire(1.0))
        time.sleep(0.5)
        self.assertFalse(b.try_acquire(1.0))
        self.assertGreaterEqual(b.tokens(), 0.4)


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