#!/usr/bin/env python3
from __future__ import annotations

import json
import subprocess
import sys
import tempfile
from pathlib import Path


HERE = Path(__file__).resolve().parent
CLIENT = HERE / "opprof_phase5_client.py"


def write_jsonl(path: Path, rows: list[dict]) -> None:
    path.write_text("".join(json.dumps(row) + "\n" for row in rows))


def run(command: list[str]) -> None:
    completed = subprocess.run(command, text=True, capture_output=True)
    if completed.returncode:
        raise RuntimeError(f"command failed: {completed.stderr}")


def main() -> None:
    with tempfile.TemporaryDirectory() as tmp_text:
        tmp = Path(tmp_text)
        manifest = tmp / "p3.jsonl"
        source = tmp / "source.jsonl"
        base = tmp / "base.jsonl"
        a1 = tmp / "a1.jsonl"
        rows = []
        source_rows = []
        lengths = [128, 8192, 256, 4096] * 8
        for index, length in enumerate(lengths):
            rows.append({
                "request_id": f"P10-{index}", "pattern_id": "P10",
                "input_tokens": length, "output_tokens": 8,
                "arrival": "steady", "kind": "private-trace",
                "source_index": index, "prompt": f"private-{index}",
            })
            source_rows.append({"timestamp": index * 0.1, "prompt": f"private-{index}"})
        write_jsonl(manifest, rows)
        write_jsonl(source, source_rows)
        common = [
            sys.executable, str(CLIENT), "transform", "--in", str(manifest),
            "--take-first", "32", "--timestamp-source", str(source),
            "--join-key", "source_index", "--timestamp-field", "timestamp",
            "--arrival", "recorded-scaled", "--target-rate", "0.5",
        ]
        run(common + ["--service-order", "original", "--out", str(base)])
        run(common + [
            "--service-order", "length-binned", "--reorder-block-size", "32",
            "--analysis-cohort-size", "16", "--max-added-delay-seconds", "64",
            "--out", str(a1),
        ])
        base_summary = json.loads((base.with_suffix(".jsonl.summary.json")).read_text())
        a1_summary = json.loads((a1.with_suffix(".jsonl.summary.json")).read_text())
        assert base_summary["rows"] == a1_summary["rows"] == 32
        assert base_summary["request_id_set_sha256"] == a1_summary["request_id_set_sha256"]
        assert a1_summary["r16"] < base_summary["r16"]
        assert a1_summary["max_added_delay_seconds"] <= 64
        assert "private-" not in json.dumps(a1_summary)
    print("phase5 tools: PASS")


if __name__ == "__main__":
    main()