#!/usr/bin/env python3
"""Split worker idle time across compile, near tile, far-parent and LOD2 phases."""
from __future__ import annotations

import csv
import json
import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l104/timeline"
EVENT = re.compile(
    r"^\[CpuTileTimeline\] event=(\S+) key=\(([^,]+),([^,]+),([^,]+),lod(-?\d+)\) "
    r"at_ms=([0-9.]+).*?worker_id=(-?\d+) queue_depth=(\d+) active_workers=(\d+) "
    r"priority_rank=(\d+)", re.M)
COLD = re.compile(
    r"^\[ColdTimeline\] event=(\S+) edge=(\S+) at_ms=([0-9.]+)", re.M)
PIPELINE = re.compile(r"^\[StreamPipeline\] (.*)$", re.M)
NEAR = re.compile(r"^\[WorldAssemblerNearLODReady\] lods=0-2 tiles=(\d+)/(\d+) cold_ready_s=([0-9.]+)$", re.M)
FIELDS = re.compile(r"([a-zA-Z0-9_]+)=([^ ]+)")


def run(path: Path) -> dict:
    content = path.read_text(encoding="utf-8", errors="replace")
    events = []
    for match in EVENT.finditer(content):
        event, x, y, z, lod, ms, worker, queue, active, rank = match.groups()
        events.append({"event": event, "lod": int(lod), "at_ms": float(ms),
                       "worker": int(worker), "queue": int(queue),
                       "active": int(active), "rank": int(rank)})
    cold = {f"{event}_{edge}": float(ms) for event, edge, ms in COLD.findall(content)}
    idle_start: dict[int, float] = {}
    idle_intervals = []
    for event in sorted((e for e in events if e["event"] in ("idle_begin", "idle_end")),
                        key=lambda row: row["at_ms"]):
        if event["event"] == "idle_begin":
            idle_start[event["worker"]] = event["at_ms"]
        elif event["worker"] in idle_start:
            idle_intervals.append((idle_start.pop(event["worker"]), event["at_ms"]))

    near = NEAR.search(content)
    pipeline = PIPELINE.search(content)
    config = re.search(r"^\[HeadlessConfig\] (.*)$", content, re.M)
    if not events or near is None or pipeline is None or config is None:
        raise RuntimeError(f"incomplete timeline markers: {path}")
    config_fields = dict(FIELDS.findall(config.group(1)))
    workers = int(config_fields["workers"])
    whole_fields = dict(FIELDS.findall(pipeline.group(1)))
    requested = [e for e in events if e["event"] == "requested"]
    admitted = [e for e in events if e["event"] == "admitted"]
    lod01_admitted = sum(e["lod"] <= 1 for e in admitted)
    lod01_requested = sum(e["lod"] <= 1 for e in requested)
    phases = (
        ("full_pipeline", "desired_tile_set_requested", "lod0_2_done_pipeline_drain"),
        ("base_shape_compile", "base_compile_begin", "base_compile_end"),
        ("lod0_1_tile_drain", "lod0_1_admission_end", "lod0_1_done_publish_drained"),
        ("far_parent_build", "far_parent_build_begin", "far_parent_build_end"),
        ("lod2_queue_and_drain", "lod2_queue_begin", "lod2_done_pipeline_drain"),
    )
    phase_rows = []
    for label, begin, end in phases:
        start, finish = cold[begin], cold[end]
        idle_worker_ms = sum(max(0.0, min(b, finish) - max(a, start))
                             for a, b in idle_intervals)
        worker_capacity_ms = max(0.0, finish - start) * workers
        phase_rows.append({"phase": label, "wall_ms": finish - start,
                           "idle_worker_ms": idle_worker_ms,
                           "idle_pct": 100 * idle_worker_ms / worker_capacity_ms
                               if worker_capacity_ms else 0.0,
                           "busy_pct": 100 * (1 - idle_worker_ms / worker_capacity_ms)
                               if worker_capacity_ms else 0.0})

    # Report the first and last near-tile dispatches and the final queued jobs.
    lod01_starts = [e for e in events if e["event"] == "started" and e["lod"] <= 1]
    lod01_finishes = [e for e in events if e["event"] == "finished" and e["lod"] <= 1]
    near_late = [e for e in lod01_starts if e["at_ms"] >= cold["lod0_1_done_publish_drained"] - 100]
    return {
        "log": str(path.relative_to(ROOT)), "workers": workers,
        "source_stamp": re.search(r"^\[VoxelForgeSourceStamp\] sha256=(\S+)$", content, re.M).group(1),
        "requested_lod01_tiles": lod01_requested, "admitted_lod01_tiles": lod01_admitted,
        "lod01_started": len(lod01_starts), "lod01_finished": len(lod01_finishes),
        "near_tiles": int(near.group(1)), "whole_tiles": whole_fields["whole_tiles"],
        "lod01_done_ms": cold["lod0_1_done_publish_drained"],
        "lod01_last_dispatch_before_done_ms": max(
            (e["at_ms"] for e in lod01_starts), default=0.0),
        "lod01_starts_in_last_100ms": len(near_late),
        "worker_idle_phases": phase_rows,
    }


def main() -> int:
    rows = [run(path) for path in sorted(OUT.glob("*-w*-direct-measured-r1.log"))]
    if len(rows) != 4:
        raise SystemExit(f"expected FI/TN W8/W12 timeline logs, found {len(rows)}")
    (OUT / "analysis.json").write_text(json.dumps(rows, indent=2) + "\n", encoding="utf-8")
    with (OUT / "utilization.csv").open("w", newline="", encoding="utf-8") as stream:
        fields = ["log", "workers", "requested_lod01_tiles", "admitted_lod01_tiles",
                  "phase", "wall_ms", "idle_worker_ms", "idle_pct", "busy_pct"]
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        for row in rows:
            for phase in row["worker_idle_phases"]:
                writer.writerow({"log": row["log"], "workers": row["workers"],
                                 "requested_lod01_tiles": row["requested_lod01_tiles"],
                                 "admitted_lod01_tiles": row["admitted_lod01_tiles"], **phase})
    print("[L104TimelineAnalysis] " + json.dumps(rows, sort_keys=True), flush=True)
    return 0


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