"""Tests for Lab 01 — serving economics, KV cache and batching.

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference — must be green
"""

import importlib
import os
from dataclasses import replace

import pytest

lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))

GIB = 1024 ** 3
GB = 1000 ** 3


def model(**kwargs):
    kwargs.setdefault("name", "m")
    kwargs.setdefault("params", 7_000_000_000)
    kwargs.setdefault("layers", 32)
    kwargs.setdefault("kv_heads", 8)
    kwargs.setdefault("head_dim", 128)
    return lab.ModelShape(**kwargs)


def gpu(**kwargs):
    kwargs.setdefault("name", "g")
    kwargs.setdefault("memory_bytes", 80 * GIB)
    kwargs.setdefault("bandwidth_bytes_per_s", 3_350 * GB)
    kwargs.setdefault("flops_per_s", 989e12)
    return lab.GPU(**kwargs)


# ======================================================================================
# 1. Shapes and the KV cache
# ======================================================================================


@pytest.mark.parametrize("field", ["params", "layers", "kv_heads", "head_dim",
                                   "bytes_per_element", "bytes_per_param"])
def test_model_shape_validates(field):
    with pytest.raises(ValueError):
        model(**{field: 0})


def test_gpu_validates():
    with pytest.raises(ValueError):
        gpu(memory_bytes=0)
    with pytest.raises(ValueError):
        gpu(count=0)


def test_kv_bytes_per_token_is_the_formula():
    # 2 * 32 layers * 8 kv_heads * 128 head_dim * 2 bytes
    assert model().kv_bytes_per_token() == 2 * 32 * 8 * 128 * 2 == 131_072


def test_kv_bytes_scales_linearly_with_sequence_length():
    m = model()
    assert m.kv_bytes(1000) == 1000 * m.kv_bytes_per_token()
    assert m.kv_bytes(0) == 0
    with pytest.raises(ValueError):
        m.kv_bytes(-1)


def test_gqa_is_the_biggest_lever_on_kv_size():
    mha = model(kv_heads=64)
    gqa = model(kv_heads=8)
    mqa = model(kv_heads=1)
    assert mha.kv_bytes_per_token() == 8 * gqa.kv_bytes_per_token()
    assert gqa.kv_bytes_per_token() == 8 * mqa.kv_bytes_per_token()


def test_weight_bytes():
    assert model(params=7_000_000_000, bytes_per_param=2).weight_bytes() == 14_000_000_000
    assert model(params=7_000_000_000, bytes_per_param=1).weight_bytes() == 7_000_000_000


def test_tensor_parallel_group_aggregates():
    single = gpu()
    node = gpu(count=8)
    assert node.total_memory_bytes == 8 * single.total_memory_bytes
    assert node.total_bandwidth_bytes_per_s == 8 * single.total_bandwidth_bytes_per_s
    assert node.total_flops_per_s == pytest.approx(8 * single.total_flops_per_s)


def test_a_model_that_does_not_fit_has_no_kv_budget():
    big = model(params=70_000_000_000)          # 130 GiB of weights
    assert lab.usable_kv_bytes(gpu(), big) == 0
    assert lab.usable_kv_bytes(gpu(count=8), big) > 0


def test_usable_kv_reserves_working_space():
    m = model()
    free = gpu().total_memory_bytes - m.weight_bytes()
    assert lab.usable_kv_bytes(gpu(), m, overhead_fraction=0.0) == free
    assert lab.usable_kv_bytes(gpu(), m, overhead_fraction=0.10) == int(free * 0.9)


def test_overhead_fraction_is_validated():
    with pytest.raises(ValueError):
        lab.usable_kv_bytes(gpu(), model(), overhead_fraction=1.0)
    with pytest.raises(ValueError):
        lab.usable_kv_bytes(gpu(), model(), overhead_fraction=-0.1)


def test_max_concurrency_is_memory_over_kv():
    m, g = model(), gpu()
    expected = lab.usable_kv_bytes(g, m) // m.kv_bytes(8192)
    assert lab.max_concurrent_sequences(g, m, 8192) == expected


def test_longer_contexts_reduce_concurrency():
    m, g = model(), gpu()
    assert (lab.max_concurrent_sequences(g, m, 2048)
            > lab.max_concurrent_sequences(g, m, 8192)
            > lab.max_concurrent_sequences(g, m, 32768))


def test_zero_length_sequence_is_an_error():
    with pytest.raises(ValueError):
        lab.max_concurrent_sequences(gpu(), model(), 0)


# ======================================================================================
# 2. Prefill and decode
# ======================================================================================


def test_prefill_scales_with_tokens_and_params():
    m, g = model(), gpu()
    assert lab.prefill_ms(m, g, 2000) == pytest.approx(2 * lab.prefill_ms(m, g, 1000))
    assert lab.prefill_ms(m, g, 0) == 0.0
    with pytest.raises(ValueError):
        lab.prefill_ms(m, g, -1)


def test_prefill_is_faster_on_more_gpus():
    m = model()
    assert lab.prefill_ms(m, gpu(count=8), 2000) < lab.prefill_ms(m, gpu(), 2000)


def test_decode_per_token_falls_as_the_batch_grows():
    m, g = model(), gpu()
    times = [lab.decode_ms_per_token(m, g, batch_size=b, context_tokens=2000)
             for b in (1, 8, 32, 128)]
    assert times == sorted(times, reverse=True)


def test_batching_amortizes_the_weight_read():
    """With no KV at all, per-token time should be exactly inversely proportional."""
    m, g = model(), gpu()
    one = lab.decode_ms_per_token(m, g, batch_size=1, context_tokens=0)
    eight = lab.decode_ms_per_token(m, g, batch_size=8, context_tokens=0)
    assert eight == pytest.approx(one / 8)


def test_the_kv_read_does_not_amortize():
    """With a large context, the KV term dominates and batching stops helping."""
    m, g = model(), gpu()
    one = lab.decode_ms_per_token(m, g, batch_size=1, context_tokens=200_000)
    eight = lab.decode_ms_per_token(m, g, batch_size=8, context_tokens=200_000)
    assert eight > one / 8            # nowhere near an 8x improvement


def test_decode_validates():
    with pytest.raises(ValueError):
        lab.decode_ms_per_token(model(), gpu(), batch_size=0)
    with pytest.raises(ValueError):
        lab.decode_ms_per_token(model(), gpu(), context_tokens=-1)


def test_decode_at_batch_one_is_far_below_the_ridge_point():
    m, g = model(), gpu()
    assert lab.arithmetic_intensity(m, 1, 2000) < lab.ridge_point(g)


def test_arithmetic_intensity_rises_with_batch_size():
    m = model()
    assert (lab.arithmetic_intensity(m, 1, 2000)
            < lab.arithmetic_intensity(m, 32, 2000))


def test_arithmetic_intensity_validates():
    with pytest.raises(ValueError):
        lab.arithmetic_intensity(model(), 0, 0)


# ======================================================================================
# 3. Batching simulators
# ======================================================================================


def test_serving_request_validates():
    with pytest.raises(ValueError):
        lab.ServingRequest("r", -1, 10, 10)
    with pytest.raises(ValueError):
        lab.ServingRequest("r", 0, 0, 10)
    with pytest.raises(ValueError):
        lab.ServingRequest("r", 0, 10, 0)


def node():
    return gpu(count=8)


def big_model():
    return model(params=70_000_000_000, layers=80)


def uneven_workload(n=24):
    return [lab.ServingRequest(f"r{i:02d}", arrival_tick=i // 4, prompt_tokens=512,
                               output_tokens=400 if i % 8 == 0 else 20)
            for i in range(n)]


def test_continuous_batching_completes_everything():
    report = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(uneven_workload())
    assert report.admitted == 24
    assert report.rejected == []


def test_continuous_beats_static_on_uneven_outputs():
    work = uneven_workload()
    cont = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(work)
    static = lab.StaticBatcher(big_model(), node(), max_batch=8).run(work)
    assert cont.total_ticks < static.total_ticks
    assert cont.mean_latency_ticks() < static.mean_latency_ticks()


def test_both_batchers_serve_the_same_requests():
    work = uneven_workload()
    cont = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(work)
    static = lab.StaticBatcher(big_model(), node(), max_batch=8).run(work)
    assert ([r.request_id for r in cont.results]
            == [r.request_id for r in static.results]
            == sorted(r.request_id for r in work))


def test_static_batching_makes_everyone_wait_for_the_longest():
    work = [lab.ServingRequest("short", 0, 100, 5),
            lab.ServingRequest("long", 0, 100, 200)]
    static = lab.StaticBatcher(big_model(), node(), max_batch=8).run(work)
    finishes = {r.request_id: r.finished_tick for r in static.results}
    assert finishes["short"] == finishes["long"]


def test_continuous_retires_the_short_one_first():
    work = [lab.ServingRequest("short", 0, 100, 5),
            lab.ServingRequest("long", 0, 100, 200)]
    cont = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(work)
    finishes = {r.request_id: r.finished_tick for r in cont.results}
    assert finishes["short"] < finishes["long"]


def test_batch_is_capped():
    work = [lab.ServingRequest(f"r{i}", 0, 100, 10) for i in range(50)]
    report = lab.ContinuousBatcher(big_model(), node(), max_batch=4).run(work)
    assert report.peak_batch == 4
    assert report.admitted == 50          # queued, not dropped


def test_max_batch_must_be_positive():
    with pytest.raises(ValueError):
        lab.ContinuousBatcher(model(), gpu(), max_batch=0)


def test_a_request_that_can_never_fit_is_rejected_not_queued():
    small = gpu(memory_bytes=24 * GIB)
    report = lab.ContinuousBatcher(model(), small, max_batch=64).run([
        lab.ServingRequest("ok-1", 0, 1_000, 100),
        lab.ServingRequest("too-big", 0, 100_000, 100_000),
        lab.ServingRequest("ok-2", 0, 1_000, 100),
    ])
    assert report.rejected == ["too-big"]
    assert report.admitted == 2


def test_admission_budgets_the_final_length_not_the_prompt():
    """A sequence that fits as a prompt but not once grown must not be admitted."""
    small = gpu(memory_bytes=24 * GIB)
    batcher = lab.ContinuousBatcher(model(), small, max_batch=64)
    per_token = model().kv_bytes_per_token()
    budget_tokens = batcher.kv_budget // per_token
    # prompt alone fits; prompt + output does not
    request = lab.ServingRequest("greedy", 0, budget_tokens - 10, 100)
    report = batcher.run([request])
    assert report.rejected == ["greedy"]


def test_peak_kv_never_exceeds_the_budget():
    batcher = lab.ContinuousBatcher(big_model(), node(), max_batch=64)
    report = batcher.run(uneven_workload(40))
    assert report.peak_kv_bytes <= batcher.kv_budget


def test_late_arrivals_are_served():
    work = [lab.ServingRequest("early", 0, 100, 5),
            lab.ServingRequest("late", 50, 100, 5)]
    report = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(work)
    assert report.admitted == 2
    by_id = {r.request_id: r for r in report.results}
    assert by_id["late"].admitted_tick >= 50


def test_empty_workload_terminates():
    report = lab.ContinuousBatcher(big_model(), node()).run([])
    assert report.admitted == 0
    assert report.total_ticks == 0


def test_report_metrics():
    report = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(uneven_workload())
    assert report.throughput_per_tick() > 0
    assert report.mean_latency_ticks() > 0
    assert report.mean_ttft_ticks() >= 0


def test_simulation_is_deterministic():
    work = uneven_workload()
    a = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(work)
    b = lab.ContinuousBatcher(big_model(), node(), max_batch=8).run(work)
    assert a.results == b.results
    assert a.total_ticks == b.total_ticks


# ======================================================================================
# 4. Capacity economics
# ======================================================================================


def payg():
    return lab.PaygPricing(input_micros_per_1k=3_000, output_micros_per_1k=12_000)


def ptu():
    return lab.ProvisionedPricing(units=100, micros_per_unit_month=120_000_000,
                                  tokens_per_unit_month=40_000_000)


def test_payg_cost_uses_both_rates():
    assert payg().cost_micros(1000, 0) == 3_000
    assert payg().cost_micros(0, 1000) == 12_000
    assert payg().cost_micros(1000, 1000) == 15_000


def test_provisioned_totals():
    assert ptu().monthly_cost_micros() == 12_000_000_000
    assert ptu().monthly_capacity_tokens() == 4_000_000_000


def test_break_even_moves_with_the_output_mix():
    """Output costs 4x input here, so a heavier output mix reaches break-even sooner."""
    low = lab.break_even(ptu(), payg(), output_fraction=0.1)
    high = lab.break_even(ptu(), payg(), output_fraction=0.5)
    assert high.break_even_tokens < low.break_even_tokens
    assert high.break_even_utilization < low.break_even_utilization


def test_break_even_arithmetic():
    be = lab.break_even(ptu(), payg(), output_fraction=0.25)
    blended = 3_000 * 0.75 + 12_000 * 0.25          # 5250 micro-USD per 1k
    assert be.break_even_tokens == int(12_000_000_000 / blended * 1000)


def test_break_even_validates():
    with pytest.raises(ValueError):
        lab.break_even(ptu(), payg(), output_fraction=1.5)


def test_spillover_fills_the_floor_first():
    plan = lab.plan_with_spillover(2_000_000_000, ptu(), payg(), output_fraction=0.25)
    assert plan.provisioned_tokens == 2_000_000_000
    assert plan.payg_tokens == 0
    assert plan.payg_micros == 0


def test_spillover_charges_the_overflow():
    plan = lab.plan_with_spillover(6_000_000_000, ptu(), payg(), output_fraction=0.25)
    assert plan.provisioned_tokens == 4_000_000_000
    assert plan.payg_tokens == 2_000_000_000
    assert plan.payg_micros > 0
    assert plan.total_micros == plan.provisioned_micros + plan.payg_micros


def test_spillover_validates():
    with pytest.raises(ValueError):
        lab.plan_with_spillover(-1, ptu(), payg(), output_fraction=0.25)


def test_self_hosted_cost_includes_the_engineering_line():
    bare = lab.SelfHostedPricing(gpus=8, micros_per_gpu_hour=3_000_000)
    with_people = lab.SelfHostedPricing(gpus=8, micros_per_gpu_hour=3_000_000,
                                        engineering_micros_per_month=25_000_000)
    assert with_people.monthly_cost_micros() - bare.monthly_cost_micros() == 25_000_000


def test_self_hosted_unit_cost_falls_with_utilization():
    hosted = lab.SelfHostedPricing(gpus=8, micros_per_gpu_hour=3_000_000)
    low = lab.self_hosted_cost_per_1k_tokens(hosted, 1_000_000_000)
    high = lab.self_hosted_cost_per_1k_tokens(hosted, 20_000_000_000)
    assert high < low
    with pytest.raises(ValueError):
        lab.self_hosted_cost_per_1k_tokens(hosted, 0)


def test_cheapest_option_moves_with_volume():
    hosted = lab.SelfHostedPricing(gpus=8, micros_per_gpu_hour=3_000_000,
                                   engineering_micros_per_month=25_000_000)
    small, _ = lab.cheapest_option(1_000_000_000, provisioned=ptu(), payg=payg(),
                                   self_hosted=hosted, output_fraction=0.25)
    large, _ = lab.cheapest_option(20_000_000_000, provisioned=ptu(), payg=payg(),
                                   self_hosted=hosted, output_fraction=0.25)
    assert small is lab.Capacity.PAYG
    assert large is lab.Capacity.SELF_HOSTED


def test_cheapest_option_without_self_hosting():
    choice, cost = lab.cheapest_option(100_000_000, provisioned=ptu(), payg=payg(),
                                       output_fraction=0.25)
    assert choice is lab.Capacity.PAYG
    assert cost > 0


def test_ties_break_toward_payg():
    """No commitment and no operational burden wins a tie."""
    free_ptu = lab.ProvisionedPricing(units=1, micros_per_unit_month=0,
                                      tokens_per_unit_month=0)
    zero_payg = lab.PaygPricing(0, 0)
    choice, cost = lab.cheapest_option(1_000, provisioned=free_ptu, payg=zero_payg,
                                       output_fraction=0.5)
    assert choice is lab.Capacity.PAYG
    assert cost == 0