#!/usr/bin/env python3
"""Five interleaved L80-protocol W8/W12 cold-readiness pairs per owner stratum."""
from __future__ import annotations

import json
import re
import statistics
import subprocess
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[4]
DATA = ROOT / "Standalone/data/owner_world.json"
BIN = ROOT / "Saved/Linux/l124/bin/VoxelForgeCoreHeadless-l124-clang"
OUT = ROOT / "Saved/Linux/l124/timing"
STAMP_TOOL = ROOT / "Standalone/Tools/verify/source_stamp.py"
CAMERAS = {
    "sw": (8.0, 16.0, -372.2626953125),
    "fi": (25.8369140625, 3.19921875, -737.0),
    "tn": (19.166015625, -14.2001953125, -753.3740234375),
    "vs": (-20.0625, -5.068359375, -900.353515625),
}
FIELDS = re.compile(r"([a-zA-Z0-9_+.-]+)=([^ ]+)")
STAMP = re.compile(r"^\[VoxelForgeSourceStamp\] sha256=([0-9a-f]{64})$", re.M)
SHAPE_FREE = re.compile(
    r"^\[ShapeReadiness\] shape_free_desired=(\d+)/(\d+) near_ring=(\d+)/27",
    re.M,
)
PROFILE = re.compile(r"^\[ShapeProfile\] (.*)$", re.M)
STAGE = re.compile(r"^\[ShapeStage\] (.*)$", re.M)
MILESTONE = re.compile(r"^\[PipelineMilestone\] (.*)$", re.M)
NEAR = re.compile(
    r"^\[WorldAssemblerNearLODReady\] lods=0-2 tiles=(\d+)/(\d+) cold_ready_s=([0-9.]+)$",
    re.M,
)
PIPELINE = re.compile(r"^\[StreamPipeline\] (.*)$", re.M)


def run(stratum: str, workers: int, phase: str, replicate: int) -> dict:
    OUT.mkdir(parents=True, exist_ok=True)
    label = f"{stratum}-w{workers}-{phase}{replicate}"
    log = OUT / f"{label}.log"
    camera = ",".join(f"{value:.12f}" for value in CAMERAS[stratum]) + ",0,0"
    command = [
        "taskset", "-c", "0-11", "prlimit", "--as=6442450944", "--",
        "timeout", "1740s", str(BIN), str(DATA), "--world-assembler-prototype",
        "--stream-benchmark", "--stream-pipeline", f"--workers={workers}",
        "--clip-radius=3", "--max-clip-level=8", "--max-level=8",
        f"--camera={camera}", "--profile-lod-work", "--stream-scheduler-policy=3",
    ]
    started = time.monotonic()
    with log.open("w", encoding="utf-8") as output:
        result = subprocess.run(command, cwd=ROOT, stdout=output,
                                stderr=subprocess.STDOUT, check=False, timeout=1740)
    content = log.read_text(encoding="utf-8", errors="replace")
    near, pipeline_match = NEAR.search(content), PIPELINE.search(content)
    profile, stage, free = PROFILE.search(content), STAGE.search(content), SHAPE_FREE.search(content)
    stamp = STAMP.search(content)
    milestones = dict(FIELDS.findall(MILESTONE.search(content).group(1))) if MILESTONE.search(content) else {}
    if result.returncode or near is None or pipeline_match is None or not all((profile, stage, free, stamp)):
        raise RuntimeError(f"baseline run incomplete: {log}; exit={result.returncode}")
    if stamp.group(1) != EXPECTED_STAMP:
        raise RuntimeError(f"source stamp mismatch: {log}; got={stamp.group(1)} expected={EXPECTED_STAMP}")
    pipeline = dict(FIELDS.findall(pipeline_match.group(1)))
    profile_values = dict(FIELDS.findall(profile.group(1)))
    stage_values = dict(FIELDS.findall(stage.group(1)))
    row = {
        "stratum": stratum,
        "workers": workers,
        "phase": phase,
        "replicate": replicate,
        "exit_code": result.returncode,
        "wall_s": round(time.monotonic() - started, 3),
        "source_stamp": stamp.group(1),
        "shape_free_tiles": int(free.group(1)),
        "desired_tiles": int(free.group(2)),
        "shape_free_near_ring": int(free.group(3)),
        "shape_compile_s": float(pipeline["shape"].removesuffix("s")),
        "first_tile_ready_s": float(pipeline["first_tile_ready"].removesuffix("s")),
        "near_ready_s": float(near.group(3)),
        "near_ready_tiles": int(near.group(1)),
        "near_desired_tiles": int(near.group(2)),
        "worker_shape_stages_ms": {
            key: float(value.removesuffix("ms"))
            for key, value in profile_values.items() if key.endswith("_worker")
        },
        "shape_stage_wall_ms": {
            key: float(value)
            for key, value in stage_values.items()
            if key in {"local-shapes", "intervals+signs", "surface-band", "far-parents", "total"}
        },
        "milestones": milestones,
        "log": str(log.relative_to(ROOT)),
    }
    print("[L124HeadlessRun] " + json.dumps(row, sort_keys=True), flush=True)
    return row


def main() -> int:
    global EXPECTED_STAMP
    expected = subprocess.run(
        ["python3", "-B", str(STAMP_TOOL), "stamp", "--root", str(ROOT)],
        cwd=ROOT, text=True, stdout=subprocess.PIPE, check=True,
    ).stdout.strip()
    EXPECTED_STAMP = expected
    if not DATA.is_file() or not BIN.is_file():
        raise SystemExit(f"missing input: {DATA} or {BIN}")
    rows: list[dict] = []
    for stratum in CAMERAS:
        for workers in (8, 12):
            run(stratum, workers, "warmup", 0)
    orders = ((8, 12), (12, 8), (8, 12), (12, 8), (8, 12))
    for stratum in CAMERAS:
        for replicate, order in enumerate(orders, 1):
            paired = {}
            for workers in order:
                row = run(stratum, workers, "measured", replicate)
                rows.append(row)
                paired[workers] = row
                (OUT / "runs-partial.json").write_text(
                    json.dumps(rows, indent=2) + "\n", encoding="utf-8")
            print("[L124HeadlessPair] " + json.dumps({
                "stratum": stratum,
                "replicate": replicate,
                "w8_near_ready_s": paired[8]["near_ready_s"],
                "w12_near_ready_s": paired[12]["near_ready_s"],
            }, sort_keys=True), flush=True)
    summary = []
    for stratum in CAMERAS:
        entry = {"stratum": stratum, "interleaved_pairs": 5}
        for workers in (8, 12):
            subset = [r for r in rows if r["stratum"] == stratum and r["workers"] == workers]
            entry[f"w{workers}_near_ready_s"] = [r["near_ready_s"] for r in subset]
            entry[f"w{workers}_near_ready_median_s"] = statistics.median(
                r["near_ready_s"] for r in subset)
            entry[f"w{workers}_shape_s"] = [r["shape_compile_s"] for r in subset]
            entry[f"w{workers}_shape_median_s"] = statistics.median(
                r["shape_compile_s"] for r in subset)
        summary.append(entry)
    (OUT / "summary.json").write_text(
        json.dumps({"source_stamp": EXPECTED_STAMP, "runs": rows, "summary": summary}, indent=2) + "\n",
        encoding="utf-8",
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())