#!/usr/bin/env python3
"""Capture comparable FI/TN L106 and L107 shape-build timelines."""
from __future__ import annotations

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

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l107/timeline"
DATA = ROOT / "Standalone/data/owner_world.json"
BINARIES = {
    "before": ROOT / "Saved/Linux/l106/bin/VoxelForgeCoreHeadless-l106-clang",
    "after": ROOT / "Saved/Linux/l107/bin/VoxelForgeCoreHeadless-l107-clang",
}
STRATA = {
    "fi": {"camera": (25.8369140625, 3.19921875, -737.0), "near": 1216, "whole": 2249},
    "tn": {"camera": (19.166015625, -14.2001953125, -753.3740234375), "near": 1344, "whole": 2233},
}
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)
CONFIG = re.compile(r"^\[HeadlessConfig\] (.*)$", re.M)
STAMP = re.compile(r"^\[VoxelForgeSourceStamp\] sha256=([0-9a-f]{64})$", re.M)
FIELDS = re.compile(r"([a-zA-Z0-9_]+)=([^ ]+)")


def run(stratum: str, workers: int, build: str) -> dict:
    OUT.mkdir(parents=True, exist_ok=True)
    log = OUT / f"{build}-{stratum}-w{workers}.log"
    camera = ",".join(f"{value:.12f}" for value in STRATA[stratum]["camera"]) + ",0,0"
    command = [
        "taskset", "-c", "0-11", "prlimit", "--as=6442450944", "--", "timeout", "1740s",
        str(BINARIES[build]), 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",
        "--cold-worker-timeline",
    ]
    started = time.monotonic()
    with log.open("w", encoding="utf-8") as stream:
        result = subprocess.run(command, cwd=ROOT, stdout=stream,
                                stderr=subprocess.STDOUT, check=False, timeout=1740)
    content = log.read_text(encoding="utf-8", errors="replace")
    near, pipeline = NEAR.search(content), PIPELINE.search(content)
    config, stamp = CONFIG.search(content), STAMP.search(content)
    shape = re.search(r"^\[ShapeBreakdown\] (.*)$", content, re.M)
    profile = re.search(r"^\[ShapeProfile\] (.*)$", content, re.M)
    substage = re.search(r"^\[ShapeSubstage\] (.*)$", content, re.M)
    overlap = re.search(r"^\[ShapeStartupOverlap\] (.*)$", content, re.M)
    if result.returncode or not all((near, pipeline, config, stamp, shape, profile, substage)):
        raise RuntimeError(f"timer trace failed or missing markers: {log}; exit={result.returncode}")
    config_fields = dict(FIELDS.findall(config.group(1)))
    pipeline_fields = dict(FIELDS.findall(pipeline.group(1)))
    expected = STRATA[stratum]
    if int(config_fields["workers"]) != workers:
        raise RuntimeError(f"worker mismatch: {log}")
    if (int(near.group(1)), int(near.group(2))) != (expected["near"], expected["near"]):
        raise RuntimeError(f"incomplete near tile set: {log}")
    if tuple(map(int, pipeline_fields["whole_tiles"].split("/"))) != (expected["whole"], expected["whole"]):
        raise RuntimeError(f"incomplete whole tile set: {log}")
    row = {
        "build": build, "stratum": stratum, "workers": workers,
        "source_stamp": stamp.group(1),
        "near_ready_s": float(near.group(3)),
        "whole_view_ready_s": float(pipeline_fields["whole_view_ready"].removesuffix("s")),
        "base_shape_s": float(pipeline_fields["shape"].removesuffix("s")),
        "shape_breakdown": dict(FIELDS.findall(shape.group(1))),
        "shape_profile": dict(FIELDS.findall(profile.group(1))),
        "shape_substage": dict(FIELDS.findall(substage.group(1))),
        "shape_startup_overlap": dict(FIELDS.findall(overlap.group(1))) if overlap else None,
        "log": str(log.relative_to(ROOT)),
        "runner_wall_s": round(time.monotonic() - started, 3),
    }
    print("[L107TimerTrace] " + json.dumps(row, sort_keys=True), flush=True)
    return row


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--build", choices=("before", "after", "both"), default="both")
    args = parser.parse_args()
    builds = ("before", "after") if args.build == "both" else (args.build,)
    for build in builds:
        binary = BINARIES[build]
        if not binary.is_file():
            raise SystemExit(f"missing input binary: {binary}")
    rows = [run(stratum, workers, build)
            for stratum in ("fi", "tn") for workers in (8, 12) for build in builds]
    (OUT / f"runs-{args.build}.json").write_text(
        json.dumps(rows, indent=2) + "\n", encoding="utf-8")
    return 0


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