#!/usr/bin/env python3
"""Estimate L77 passage/support worker shares before and after L109."""
from __future__ import annotations

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

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l109/l77-phase-census"
DATA = ROOT / "Standalone/data/owner_world.json"
BINARIES = {
    "base": ROOT / "Saved/Linux/l109/bin-baseline-profile/VoxelForgeCoreHeadless-l109-base-profile",
    "candidate": ROOT / "Saved/Linux/l109/bin-final-profile/VoxelForgeCoreHeadless-l109-candidate-profile",
}
CAMERAS = {
    "FI": (25.8369140625, 3.19921875, -737.0),
    "TN": (19.166015625, -14.2001953125, -753.3740234375),
}
WORKER = re.compile(r"^\[WorldAssemblerLODWorker\] (.*)$", re.M)
SAMPLE = re.compile(r"^\[L77Sample\] (.*)$", re.M)
SAMPLER = re.compile(r"^\[L77Sampler\] action=stop (.*)$", re.M)
FIELDS = re.compile(r"([a-zA-Z0-9_]+)=([^ ]+)")


def run(stratum: str, variant: str, replicate: int) -> dict:
    OUT.mkdir(parents=True, exist_ok=True)
    label = f"{stratum.lower()}-{variant}-r{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(BINARIES[variant]), str(DATA),
        "--world-assembler-prototype", "--stream-benchmark", "--stream-pipeline",
        "--workers=8", "--clip-radius=3", "--max-clip-level=8", "--max-level=8",
        f"--camera={camera}", "--profile-lod-work", "--stream-scheduler-policy=3",
        "--linux77-profile",
    ]
    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")
    sampler_match = SAMPLER.search(content)
    workers = [dict(FIELDS.findall(match.group(1))) for match in WORKER.finditer(content)]
    samples = [dict(FIELDS.findall(match.group(1))) for match in SAMPLE.finditer(content)]
    if result.returncode or sampler_match is None or len(workers) < 9 or not samples:
        raise SystemExit(f"phase census failed or markers missing: {log}; exit={result.returncode}")
    sampler = dict(FIELDS.findall(sampler_match.group(1)))
    if sampler.get("status") != "PASS":
        raise SystemExit(f"L77 sampler failed: {log}: {sampler}")
    lod_seconds = {int(row["lod"]): float(row["worker_sum_s"]) for row in workers}

    def estimate(scope: str) -> dict:
        lods = set(range(3)) if scope == "lod0-2" else set(range(9))
        eligible = [row for row in samples if int(row["lod"]) in lods]
        weights = {phase: sum(int(row.get("weight", "1")) for row in eligible
                              if row.get("phase") == phase)
                   for phase in ("passages", "support")}
        total_weight = sum(int(row.get("weight", "1")) for row in eligible)
        cpu_seconds = sum(value for lod, value in lod_seconds.items() if lod in lods)
        result_row = {"worker_cpu_s": round(cpu_seconds, 4),
                      "weighted_samples": total_weight}
        for phase, weight in weights.items():
            share = weight / total_weight if total_weight else 0.0
            result_row[f"{phase}_share_pct"] = round(share * 100.0, 3)
            result_row[f"{phase}_worker_s_estimate"] = round(cpu_seconds * share, 4)
        return result_row

    row = {
        "stratum": stratum, "variant": variant, "replicate": replicate,
        "exit_code": result.returncode,
        "sampler": sampler,
        "lod0_2": estimate("lod0-2"), "whole_view": estimate("whole"),
        "log": str(log.relative_to(ROOT)),
        "runner_wall_s": round(time.monotonic() - started, 3),
    }
    print("[L109PhaseCensusRun] " + json.dumps(row, sort_keys=True), flush=True)
    return row


def main() -> int:
    if any(not path.is_file() for path in BINARIES.values()):
        raise SystemExit("missing L77 profile binary")
    rows = [run(stratum, variant, replicate)
            for stratum in ("FI", "TN")
            for replicate in (1, 2)
            for variant in ("base", "candidate")]
    summary = []
    for stratum in ("FI", "TN"):
        for variant in ("base", "candidate"):
            entry = {"stratum": stratum, "variant": variant, "replicates": 2}
            group = [row for row in rows if row["stratum"] == stratum
                     and row["variant"] == variant]
            for scope in ("lod0_2", "whole_view"):
                for key in ("worker_cpu_s", "passages_share_pct", "passages_worker_s_estimate",
                            "support_share_pct", "support_worker_s_estimate"):
                    values = [row[scope][key] for row in group]
                    entry[f"{scope}_{key}_runs"] = values
                    entry[f"{scope}_{key}_median"] = statistics.median(values)
            summary.append(entry)
    (OUT / "summary.json").write_text(
        json.dumps({"runs": rows, "summary": summary}, indent=2) + "\n",
        encoding="utf-8")
    return 0


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