#!/usr/bin/env python3
"""Symbolize L137 L77 samples and compare inclusive worker functions."""
from __future__ import annotations

import csv
import json
import re
import subprocess
from collections import Counter
from pathlib import Path

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l137/profile"
BINARIES = {
    "assembler-world": ROOT / "Saved/Linux/l137/profile-bin/VoxelForgeCoreHeadless-l137-profile-base",
    "head": ROOT / "Saved/Linux/l137/profile-bin/VoxelForgeCoreHeadless-l137-profile-head",
}
SAMPLE_RE = re.compile(
    r"^\[L77Sample\] pc=(0x\w+) lod=(\d+) stage=(\S+) route=(\S+) "
    r"phase=(\S+) halo=(\d+) weight=(\d+)$", re.M)


def samples(version: str) -> tuple[list[dict], set[str]]:
    rows, pcs = [], set()
    for replicate in (1, 2, 3):
        path = OUT / f"sw-{version}-r{replicate}.log"
        for pc, lod, stage, route, phase, halo, weight in SAMPLE_RE.findall(
                path.read_text(encoding="utf-8", errors="replace")):
            lod_value = int(lod)
            if not 0 <= lod_value <= 8:
                continue
            pcs.add(pc)
            rows.append({"version": version, "replicate": replicate,
                         "pc": pc, "lod": lod_value, "stage": stage,
                         "route": route, "phase": phase, "halo": int(halo),
                         "weight": int(weight)})
    return rows, pcs


def symbolize(version: str, pcs: set[str]) -> dict[str, str]:
    command = ["llvm-symbolizer-19", "--demangle", "--inlines", "--output-style=JSON",
               f"--obj={BINARIES[version]}"]
    result = subprocess.run(command, input="\n".join(sorted(pcs)) + "\n",
                            capture_output=True, text=True, check=True)
    decoder = json.JSONDecoder()
    offset, output = 0, {}
    while offset < len(result.stdout):
        while offset < len(result.stdout) and result.stdout[offset].isspace():
            offset += 1
        if offset == len(result.stdout):
            break
        item, offset = decoder.raw_decode(result.stdout, offset)
        frames = item.get("Symbol", [])
        output[item["Address"]] = (
            frames[0].get("FunctionName", "?") if frames else "?")
    if output.keys() != pcs:
        raise RuntimeError(f"{version}: symbolizer resolved {len(output)}/{len(pcs)} PCs")
    return output


def main() -> int:
    totals: dict[str, Counter] = {}
    lod_totals: dict[str, Counter] = {}
    detail_rows = []
    for version in BINARIES:
        rows, pcs = samples(version)
        symbols = symbolize(version, pcs)
        totals[version] = Counter()
        lod_totals[version] = Counter()
        for row in rows:
            function = symbols[row["pc"]]
            weight = row["weight"]
            totals[version][function] += weight
            lod_totals[version][row["lod"]] += weight
            detail_rows.append({**row, "function": function})
    total_by_version = {key: sum(counter.values()) for key, counter in totals.items()}
    functions = set(totals["assembler-world"]) | set(totals["head"])
    result = []
    for function in functions:
        base_ticks = totals["assembler-world"][function]
        head_ticks = totals["head"][function]
        base_pct = 100.0 * base_ticks / total_by_version["assembler-world"]
        head_pct = 100.0 * head_ticks / total_by_version["head"]
        result.append({
            "function": function,
            "assembler_world_ticks": base_ticks,
            "assembler_world_share_pct": round(base_pct, 3),
            "head_ticks": head_ticks,
            "head_share_pct": round(head_pct, 3),
            "share_delta_pp": round(head_pct - base_pct, 3),
        })
    result.sort(key=lambda row: abs(row["share_delta_pp"]), reverse=True)
    with (OUT / "function-diff.csv").open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(result[0]))
        writer.writeheader()
        writer.writerows(result)
    with (OUT / "symbolized-samples.csv").open(
            "w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(detail_rows[0]))
        writer.writeheader()
        writer.writerows(detail_rows)
    summary = {
        "scope": "worker samples LOD0-8, all LODs included; lod=255 main-thread samples excluded",
        "runs_per_version": 3,
        "weighted_ticks": total_by_version,
        "lod_weighted_ticks": {version: dict(sorted(values.items()))
                               for version, values in lod_totals.items()},
        "function_diff": result,
    }
    (OUT / "function-diff.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    print("[L137ProfileDiff] " + json.dumps({
        "scope": summary["scope"], "weighted_ticks": total_by_version,
        "lod_weighted_ticks": summary["lod_weighted_ticks"],
        "largest_share_deltas": result[:20]}, sort_keys=True))
    return 0


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