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


SCRIPT = (
    Path(__file__).resolve().parents[1]
    / "runs/frontier-workload-regime-taxonomy-v0/project_per_gpu_load.py"
)
SPEC = importlib.util.spec_from_file_location("project_per_gpu_load", 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 test_project_case_scales_only_arrivals_and_load_metadata(tmp_path):
    public_path = tmp_path / "source.csv"
    with public_path.open("w", newline="") as output:
        writer = csv.DictWriter(
            output,
            fieldnames=("arrived_at", "num_prefill_tokens", "num_decode_tokens"),
        )
        writer.writeheader()
        writer.writerow(
            {"arrived_at": "0", "num_prefill_tokens": "16", "num_decode_tokens": "8"}
        )
        writer.writerow(
            {"arrived_at": "4", "num_prefill_tokens": "32", "num_decode_tokens": "12"}
        )
    private_path = tmp_path / "source.jsonl"
    private_path.write_text(
        "\n".join(
            json.dumps({"arrived_at": arrival, "body": {"prompt": index}})
            for index, arrival in enumerate((0.0, 4.0))
        )
        + "\n"
    )
    case = {
        "family": "fixture",
        "rho": 0.01,
        "global_offered_request_rate": 0.25,
        "empirical_interarrival_rate": 0.25,
        "decode_offered_tokens_per_second": 2.5,
        "last_arrival_s": 4.0,
        "public_csv": str(public_path),
        "public_csv_sha256": MODULE.sha256(public_path),
        "private_jsonl": str(private_path),
        "private_jsonl_sha256": MODULE.sha256(private_path),
    }

    projected = MODULE.project_case(case, tmp_path / "output", 4)

    with Path(projected["public_csv"]).open(newline="") as source:
        public_rows = list(csv.DictReader(source))
    private_rows = [
        json.loads(line) for line in Path(projected["private_jsonl"]).open()
    ]
    assert [float(row["arrived_at"]) for row in public_rows] == [0.0, 1.0]
    assert [row["arrived_at"] for row in private_rows] == [0.0, 1.0]
    assert [row["num_decode_tokens"] for row in public_rows] == ["8", "12"]
    assert projected["per_gpu_offered_request_rate"] == 0.25
    assert projected["global_offered_request_rate"] == 1.0
    assert projected["decode_offered_tokens_per_second"] == 10.0
    assert projected["last_arrival_s"] == 1.0