#!/usr/bin/env python3
"""Attribute FI/TN L77 sample PCs to current-source functions and work areas."""
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/l104/census"
BINARY = ROOT / "Saved/Linux/l104/profile-bin/VoxelForgeCoreHeadless-l104-profile-clang"
SAMPLE = re.compile(
    r"^\[L77Sample\] pc=(0x\w+) lod=(\d+) stage=(\S+) route=(\S+) "
    r"phase=(\S+) halo=(\d+) weight=(\d+)$", re.M)
TARGETS = {
    "sweep_base_prepare": "prepareSweptTubeBase",
    "sweep_lookup": "sweptSample",
    "exact_route_visit": "VisitNearestPassage",
    "miss_reach_bound": "PassagePotentialWithinNode",
    "far_shell_down_hit": "HasDownwardHit",
    "far_shell_up_hit": "HasUpwardHit",
}


def load_samples() -> tuple[list[dict], set[str]]:
    rows, pcs = [], set()
    for stratum in ("FI", "TN"):
        for replicate in (1, 2, 3):
            log = OUT / f"{stratum.lower()}-r{replicate}.log"
            content = log.read_text(encoding="utf-8", errors="replace")
            for pc, lod, stage, route, phase, halo, weight in SAMPLE.findall(content):
                if int(lod) > 2:
                    continue
                pcs.add(pc)
                rows.append({"stratum": stratum, "replicate": replicate, "pc": pc,
                             "lod": int(lod), "stage": stage, "route": route,
                             "phase": phase, "weight": int(weight)})
    return rows, pcs


def symbolize(pcs: set[str]) -> dict[str, list[dict]]:
    command = ["llvm-symbolizer-19", "--demangle", "--inlines", "--output-style=JSON",
               f"--obj={BINARY}"]
    result = subprocess.run(command, input="\n".join(sorted(pcs)) + "\n",
                            capture_output=True, text=True, check=True)
    decoder = json.JSONDecoder()
    offset = 0
    output = {}
    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)
        output[item["Address"]] = item.get("Symbol", [])
    if output.keys() != pcs:
        raise RuntimeError(f"symbolizer resolved {len(output)}/{len(pcs)} sampled PCs")
    return output


def write_csv(path: Path, rows: list[dict]) -> None:
    with path.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def main() -> int:
    if not BINARY.is_file():
        raise SystemExit(f"missing profile binary: {BINARY}")
    samples, pcs = load_samples()
    symbols = symbolize(pcs)
    function_totals = Counter()
    target_totals = Counter()
    for sample in samples:
        frames = symbols[sample["pc"]]
        top_function = frames[0].get("FunctionName", "?") if frames else "?"
        function_totals[(sample["stratum"], sample["stage"], sample["phase"],
                         top_function)] += sample["weight"]
        seen = set()
        for name, token in TARGETS.items():
            matches = [frame for frame in frames if token in frame.get("FunctionName", "")]
            if not matches:
                continue
            frame = matches[0]
            line = int(frame.get("Line", 0))
            category = name
            if name == "sweep_base_prepare":
                category = ("sweep_prepare_base" if 14249 <= line <= 14435 else
                            "sweep_prepare_detail" if 14436 <= line <= 14521 else
                            "sweep_prepare_unlocated")
            key = (sample["stratum"], sample["stage"], sample["phase"], category)
            if key not in seen:
                target_totals[key] += sample["weight"]
                seen.add(key)

    by_function = [
        {"stratum": stratum, "stage": stage, "phase": phase,
         "weighted_ticks": ticks, "top_function": function}
        for (stratum, stage, phase, function), ticks in
        sorted(function_totals.items(), key=lambda item: (item[0][:3], -item[1]))
    ]
    by_target = [
        {"stratum": stratum, "stage": stage, "phase": phase,
         "work_area": target, "weighted_ticks": ticks}
        for (stratum, stage, phase, target), ticks in sorted(target_totals.items())
    ]
    write_csv(OUT / "symbol-function-totals.csv", by_function)
    write_csv(OUT / "symbol-work-area-totals.csv", by_target)
    print(f"[L103SymbolAnalysis] samples={len(samples)} unique_pcs={len(pcs)} "
          f"functions={len(by_function)} work_areas={len(by_target)}")
    return 0


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