import importlib.util
import json
import sys
from pathlib import Path


SCRIPT = (
    Path(__file__).resolve().parents[1]
    / "runs/frontier-workload-regime-taxonomy-v0/materialize_workload_regimes.py"
)
SPEC = importlib.util.spec_from_file_location("workload_regimes", SCRIPT)
MODULE = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
sys.modules[SPEC.name] = MODULE
SPEC.loader.exec_module(MODULE)


def source_rows():
    return [
        {
            "source_index": index,
            "source_arrived_at": arrival,
            "input_length": input_length,
            "output_length": output_length,
            "session_id": 100 + index,
            "runtime_block_ids": list(range(input_length // 16)),
            "body": {"prompt": f"prompt-{index}", "model": "fixture"},
        }
        for index, (arrival, input_length, output_length) in enumerate(
            [(10.0, 16, 4), (10.5, 32, 8), (14.0, 48, 12)]
        )
    ]


def test_arrival_contracts_preserve_rate_and_burst_pattern():
    rows = source_rows()
    uniform = MODULE.arrivals_for(rows, "uniform", 2.0)
    traced = MODULE.arrivals_for(rows, "trace", 2.0)

    assert uniform == [0.0, 0.5, 1.0]
    assert (len(traced) - 1) / (traced[-1] - traced[0]) == 2.0
    assert traced[1] - traced[0] < traced[2] - traced[1]


def test_family_axes_change_only_the_declared_contract():
    rows = source_rows()
    token_ids = [10, 11, 12, 13]
    families = {family.name: family for family in MODULE.FAMILIES}

    w1 = MODULE.build_family_rows(
        rows,
        families["w1-mean-fixed-uniform-none"],
        target_rate=1.0,
        fixed_token_ids=token_ids,
    )
    w3 = MODULE.build_family_rows(
        rows,
        families["w3-heterogeneous-uniform-none"],
        target_rate=1.0,
        fixed_token_ids=token_ids,
    )
    w5 = MODULE.build_family_rows(
        rows,
        families["w5-heterogeneous-uniform-prefix"],
        target_rate=1.0,
        fixed_token_ids=token_ids,
    )

    assert {(row["input_length"], row["output_length"]) for row in w1} == {(32, 8)}
    assert [(row["input_length"], row["output_length"]) for row in w3] == [
        (16, 4),
        (32, 8),
        (48, 12),
    ]
    assert all(row["runtime_block_ids"] == [] for row in w1 + w3)
    assert [row["runtime_block_ids"] for row in w5] == [
        [0],
        [0, 1],
        [0, 1, 2],
    ]
    assert [row["body"]["prompt"] for row in w3] == [
        "prompt-0",
        "prompt-1",
        "prompt-2",
    ]


def test_private_source_projects_only_complete_prefix_blocks(tmp_path):
    private_path = tmp_path / "source.jsonl"
    source = {
        "source_index": 7,
        "arrived_at": 1.5,
        "input_length": 17,
        "output_length": 4,
        "session_id": 11,
        "runtime_block_ids": [101, 102],
        "body": {"prompt": "fixture"},
    }
    private_path.write_text(json.dumps(source) + "\n")

    rows = MODULE.load_source(None, private_path)

    assert rows[0]["runtime_block_ids"] == [101]
    assert rows[0]["source_arrived_at"] == 1.5