import threading
import time

from token_bucket import TokenBucket


def test_basic_acquire_and_refill():
    b = TokenBucket(rate=10.0, capacity=5.0)
    assert b.try_acquire(5.0) is True
    assert b.try_acquire(1.0) is False
    # refill: wait ~0.15s -> 1.5 tokens
    time.sleep(0.15)
    assert b.try_acquire(1.0) is True


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


def test_acquire_blocks_and_returns():
    b = TokenBucket(rate=20.0, capacity=1.0)
    b.try_acquire(1.0)  # drain
    t0 = time.monotonic()
    ok = b.acquire(1.0, timeout=0.2)
    dt = time.monotonic() - t0
    assert ok is True
    assert dt >= 0.04  # had to wait ~0.05s for 1 token at 20/s


def test_n_exceeds_capacity_raises():
    b = TokenBucket(rate=1.0, capacity=1.0)
    try:
        b.acquire(2.0, timeout=0.01)
        raised = False
    except ValueError:
        raised = True
    assert raised
    try:
        b.try_acquire(2.0)
        raised2 = False
    except ValueError:
        raised2 = True
    assert raised2


def test_invalid_params():
    for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
        try:
            TokenBucket(**kwargs)
            assert False
        except ValueError:
            pass
    try:
        TokenBucket(rate=1, capacity=1).try_acquire(0)
        assert False
    except ValueError:
        pass


def test_concurrent_no_overconsume():
    b = TokenBucket(rate=1e-9, capacity=10.0)  # effectively no refill
    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()
    assert successes["n"] == 10


def test_acquire_timeout_false():
    b = TokenBucket(rate=1.0, capacity=1.0)
    b.try_acquire(1.0)  # drain
    # request 1 token but timeout too short for refill (need 1s)
    t0 = time.monotonic()
    ok = b.acquire(1.0, timeout=0.05)
    dt = time.monotonic() - t0
    assert ok is False
    assert dt < 0.2


def test_reset():
    b = TokenBucket(rate=1.0, capacity=2.0)
    b.try_acquire(2.0)
    assert b.try_acquire(1.0) is False
    b.reset()
    assert b.try_acquire(2.0) is True


def test_partial_never_consumes():
    b = TokenBucket(rate=1.0, capacity=2.0)
    b.try_acquire(2.0)  # drain
    # 1 token needed, only 0 available -> no consume
    assert b.try_acquire(1.0) is False
    # refill 0.5s -> 0.5 tokens, still < 1
    time.sleep(0.5)
    assert b.try_acquire(1.0) is False
    # tokens should be ~0.5, not negative
    assert b.tokens() >= 0.4