#!/usr/bin/env python3
"""
vprof_to_json.py - Parse a Source Engine VProf recording (*.vprof) and emit JSON.

The binary format is documented in VPROF_FORMAT.md. It is produced by the
engine's ``vprof_record_start`` command (class CVProfRecorder in
engine/vprof_record.cpp).

Usage:
    python vprof_to_json.py test.vprof > test.json
    python vprof_to_json.py test.vprof -o test.json --indent 2
    python vprof_to_json.py test.vprof --summary-only
    python vprof_to_json.py test.vprof --no-frames --top 50
"""

from __future__ import annotations

import argparse
import json
import statistics
import struct
import sys
from dataclasses import dataclass, field
from typing import Dict, List, Optional

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

VPROF_FILE_VERSION = 1

TOKEN_START_FRAME = 0
TOKEN_ADD_NODE = 1
TOKEN_ADD_BUDGET_GROUP = 2
TOKEN_TIMINGS = 3
TOKEN_END_OF_FRAME = 4
TOKEN_FILE_FINISHED = 5

BUDGETFLAG_CLIENT = 1 << 0
BUDGETFLAG_SERVER = 1 << 1
BUDGETFLAG_OTHER = 1 << 2
BUDGETFLAG_HIDDEN = 1 << 15

_BUDGET_FLAG_NAMES = (
    (BUDGETFLAG_CLIENT, "CLIENT"),
    (BUDGETFLAG_SERVER, "SERVER"),
    (BUDGETFLAG_OTHER, "OTHER"),
    (BUDGETFLAG_HIDDEN, "HIDDEN"),
)

# The 31 budget groups registered by CVProfile::CVProfile() in a fixed order.
# A recording may either rewrite all of them (index 0..30) or only the groups
# added after construction; either way IDs 0..30 map to these names.
DEFAULT_BUDGET_GROUPS = [
    ("Unaccounted", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("World Rendering", BUDGETFLAG_CLIENT),
    ("Displacement_Rendering", BUDGETFLAG_CLIENT),
    ("Game", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("Player", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("NPCs", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("Server Animation", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("Client_Animation", BUDGETFLAG_CLIENT),
    ("Physics", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("Static_Prop_Rendering", BUDGETFLAG_CLIENT),
    ("Other_Model_Rendering", BUDGETFLAG_CLIENT),
    ("Fast Path Model Rendering", BUDGETFLAG_CLIENT),
    ("Light_Cache", BUDGETFLAG_CLIENT),
    ("Brush_Model_Rendering", BUDGETFLAG_CLIENT),
    ("Shadow_Rendering", BUDGETFLAG_CLIENT),
    ("Detail_Prop_Rendering", BUDGETFLAG_CLIENT),
    ("Particle/Effect_Rendering", BUDGETFLAG_CLIENT),
    ("Ropes", BUDGETFLAG_CLIENT),
    ("Dynamic_Light_Rendering", BUDGETFLAG_CLIENT),
    ("Networking", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("Sound", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("VGUI", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("FileSystem", BUDGETFLAG_SERVER | BUDGETFLAG_OTHER),
    ("Prediction", BUDGETFLAG_CLIENT),
    ("Interpolation", BUDGETFLAG_CLIENT),
    ("Swap_Buffers", BUDGETFLAG_CLIENT),
    ("Occlusion", BUDGETFLAG_CLIENT),
    ("Overlays", BUDGETFLAG_CLIENT),
    ("Tools", BUDGETFLAG_CLIENT | BUDGETFLAG_OTHER),
    ("Texture_Cache", BUDGETFLAG_CLIENT),
    ("Replay", BUDGETFLAG_SERVER),
]


class VProfError(Exception):
    """Raised when a .vprof stream is malformed."""


def decode_flags(flags: int) -> List[str]:
    return [name for bit, name in _BUDGET_FLAG_NAMES if flags & bit]


# ---------------------------------------------------------------------------
# Binary reader
# ---------------------------------------------------------------------------


class Reader:
    """Small cursor over an in-memory byte buffer, little-endian."""

    __slots__ = ("data", "pos", "size")

    def __init__(self, data: bytes):
        self.data = data
        self.pos = 0
        self.size = len(data)

    @property
    def remaining(self) -> int:
        return self.size - self.pos

    def _take(self, n: int) -> bytes:
        if self.pos + n > self.size:
            raise VProfError(
                f"unexpected end of file at offset {self.pos} (wanted {n} bytes)"
            )
        chunk = self.data[self.pos : self.pos + n]
        self.pos += n
        return chunk

    def u8(self) -> int:
        return self._take(1)[0]

    def i32(self) -> int:
        return struct.unpack("<i", self._take(4))[0]

    def u16(self) -> int:
        return struct.unpack("<H", self._take(2))[0]

    def u32(self) -> int:
        return struct.unpack("<I", self._take(4))[0]

    def cstring(self, max_len: int = 4096) -> str:
        end = self.data.find(b"\x00", self.pos)
        if end < 0:
            raise VProfError(f"unterminated string at offset {self.pos}")
        if end - self.pos > max_len:
            raise VProfError(f"string too long at offset {self.pos}")
        raw = self.data[self.pos : end]
        self.pos = end + 1
        # Source writes narrow strings; Windows-1252 is a safe superset of
        # the ASCII that is normally used.
        return raw.decode("utf-8", "replace")


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------


@dataclass
class BudgetGroup:
    id: int
    name: str
    flags: int
    is_default: bool = False


@dataclass
class Node:
    id: int
    name: str
    parent_id: Optional[int]
    budget_group_id: int
    children: List[int] = field(default_factory=list)


@dataclass
class Frame:
    index: int
    tick: int
    samples: Dict[int, tuple] = field(default_factory=dict)  # node_id -> (calls, us)
    added_nodes: List[int] = field(default_factory=list)
    added_budget_groups: List[int] = field(default_factory=list)


@dataclass
class Recording:
    version: int
    root_node_id: int
    budget_groups: List[BudgetGroup]
    nodes: Dict[int, Node]
    frames: List[Frame]


# ---------------------------------------------------------------------------
# Parser
# ---------------------------------------------------------------------------


def parse_vprof(data: bytes) -> Recording:
    r = Reader(data)

    try:
        version = r.i32()
        root_node_id = r.i32()
    except VProfError as e:
        raise VProfError(f"truncated header: {e}") from e

    if version != VPROF_FILE_VERSION:
        # Do not fail outright: other revisions may be readable.
        sys.stderr.write(
            f"warning: vprof file version is {version}, "
            f"expected {VPROF_FILE_VERSION}; attempting to parse anyway\n"
        )

    # Pre-seed the default budget groups so IDs stay correct whether or not
    # the recording rewrites them.
    groups: List[BudgetGroup] = [
        BudgetGroup(id=i, name=name, flags=flags, is_default=True)
        for i, (name, flags) in enumerate(DEFAULT_BUDGET_GROUPS)
    ]
    group_by_name: Dict[str, int] = {g.name: g.id for g in groups}

    def add_budget_group(name: str, flags: int) -> tuple:
        """Return (id, was_new)."""
        existing = group_by_name.get(name)
        if existing is not None:
            groups[existing].flags = flags
            return existing, False
        gid = len(groups)
        groups.append(BudgetGroup(id=gid, name=name, flags=flags, is_default=False))
        group_by_name[name] = gid
        return gid, True

    nodes: Dict[int, Node] = {
        root_node_id: Node(
            id=root_node_id, name="Root", parent_id=None, budget_group_id=0
        )
    }
    frames: List[Frame] = []

    def read_timings(node_id: int, samples: Dict[int, tuple]) -> None:
        node = nodes[node_id]

        calls_token = r.u8()
        if calls_token == 255:
            calls = r.u16()
        else:
            calls = calls_token

        time_units = r.u16()
        if time_units == 0xFFFF:
            microseconds = r.u32() * 4
        else:
            microseconds = time_units * 4

        samples[node_id] = (calls, microseconds)

        for child_id in node.children:
            read_timings(child_id, samples)

    while r.remaining:
        token = r.u8()

        if token == TOKEN_FILE_FINISHED:
            break
        if token != TOKEN_START_FRAME:
            raise VProfError(
                f"expected Token_StartFrame at offset {r.pos - 1}, got {token}"
            )

        tick = r.i32()
        frame = Frame(index=len(frames), tick=tick)

        while True:
            inner = r.u8()

            if inner == TOKEN_END_OF_FRAME:
                break
            elif inner == TOKEN_ADD_BUDGET_GROUP:
                name = r.cstring()
                flags = r.i32()
                gid, is_new = add_budget_group(name, flags)
                if is_new:
                    frame.added_budget_groups.append(gid)
            elif inner == TOKEN_ADD_NODE:
                parent_id = r.i32()
                name = r.cstring()
                budget_group_id = r.i32()
                node_id = r.i32()

                if parent_id not in nodes:
                    raise VProfError(
                        f"AddNode references unknown parent {parent_id} "
                        f"at offset {r.pos}"
                    )
                if node_id in nodes:
                    raise VProfError(
                        f"AddNode re-declares existing node {node_id} "
                        f"at offset {r.pos}"
                    )

                nodes[node_id] = Node(
                    id=node_id,
                    name=name,
                    parent_id=parent_id,
                    budget_group_id=budget_group_id,
                )
                # CVProfNode::GetSubNode() prepends, so mirror that here.
                nodes[parent_id].children.insert(0, node_id)
                frame.added_nodes.append(node_id)
            elif inner == TOKEN_TIMINGS:
                read_timings(root_node_id, frame.samples)
            else:
                raise VProfError(
                    f"unexpected token {inner} inside frame {frame.index} "
                    f"at offset {r.pos - 1}"
                )

        frames.append(frame)

    return Recording(
        version=version,
        root_node_id=root_node_id,
        budget_groups=groups,
        nodes=nodes,
        frames=frames,
    )


# ---------------------------------------------------------------------------
# JSON conversion
# ---------------------------------------------------------------------------


def _group_name(recording: Recording, group_id: int) -> str:
    if 0 <= group_id < len(recording.budget_groups):
        return recording.budget_groups[group_id].name
    return f"<unknown group {group_id}>"


def build_summary(recording: Recording, top_n: int) -> dict:
    totals = [
        f.samples.get(recording.root_node_id, (0, 0))[1] for f in recording.frames
    ]

    # Per-node aggregates.
    agg: Dict[int, dict] = {}
    for frame in recording.frames:
        for node_id, (calls, us) in frame.samples.items():
            entry = agg.get(node_id)
            if entry is None:
                entry = agg[node_id] = {
                    "total_calls": 0,
                    "total_microseconds": 0,
                    "max_microseconds": 0,
                    "frames_present": 0,
                }
            entry["total_calls"] += calls
            entry["total_microseconds"] += us
            entry["frames_present"] += 1
            if us > entry["max_microseconds"]:
                entry["max_microseconds"] = us

    grand_total = sum(totals) or 1
    top_nodes = []
    for node_id, entry in sorted(
        agg.items(), key=lambda kv: kv[1]["total_microseconds"], reverse=True
    )[:top_n]:
        node = recording.nodes.get(node_id)
        frames_present = entry["frames_present"] or 1
        top_nodes.append(
            {
                "node_id": node_id,
                "name": node.name if node else f"<node {node_id}>",
                "budget_group": _group_name(
                    recording, node.budget_group_id if node else 0
                ),
                "total_calls": entry["total_calls"],
                "total_microseconds": entry["total_microseconds"],
                "avg_microseconds": round(
                    entry["total_microseconds"] / frames_present, 3
                ),
                "max_microseconds": entry["max_microseconds"],
                "percent_of_total": round(
                    100.0 * entry["total_microseconds"] / grand_total, 4
                ),
            }
        )

    def stat(values: List[int]) -> Optional[dict]:
        if not values:
            return None
        ordered = sorted(values)
        result = {
            "min": ordered[0],
            "max": ordered[-1],
            "mean": round(statistics.fmean(values), 3),
            "median": round(statistics.median(values), 3),
            "total": sum(values),
        }
        if len(ordered) > 1:
            idx = min(len(ordered) - 1, int(round(0.95 * (len(ordered) - 1))))
            result["p95"] = ordered[idx]
        return result

    ticks = [f.tick for f in recording.frames]
    return {
        "frame_count": len(recording.frames),
        "node_count": len(recording.nodes),
        "budget_group_count": len(recording.budget_groups),
        "tick_first": ticks[0] if ticks else None,
        "tick_last": ticks[-1] if ticks else None,
        "tick_span": (ticks[-1] - ticks[0]) if ticks else None,
        "total_recorded_microseconds": sum(totals),
        "frame_time_microseconds": stat(totals),
        "top_nodes_by_total_microseconds": top_nodes,
    }


def recording_to_dict(
    recording: Recording,
    *,
    source_name: str,
    size_bytes: int,
    include_frames: bool,
    include_nodes: bool,
    top_n: int,
) -> dict:
    out = {
        "format": "source-vprof",
        "source": {
            "file": source_name,
            "size_bytes": size_bytes,
            "version": recording.version,
            "root_node_id": recording.root_node_id,
        },
    }

    out["budget_groups"] = [
        {
            "id": g.id,
            "name": g.name,
            "flags": g.flags,
            "flags_decoded": decode_flags(g.flags),
            "is_default": g.is_default,
        }
        for g in recording.budget_groups
    ]

    if include_nodes:
        out["nodes"] = [
            {
                "id": n.id,
                "name": n.name,
                "parent_id": n.parent_id,
                "budget_group_id": n.budget_group_id,
                "budget_group": _group_name(recording, n.budget_group_id),
                "children": n.children,
            }
            for n in recording.nodes.values()
        ]

    if include_frames:
        frames_out = []
        for frame in recording.frames:
            total_us = frame.samples.get(recording.root_node_id, (0, 0))[1]
            frames_out.append(
                {
                    "index": frame.index,
                    "tick": frame.tick,
                    "total_microseconds": total_us,
                    # node_id (string key) -> {"calls": int, "microseconds": int}
                    "samples": {
                        str(node_id): {
                            "calls": calls,
                            "microseconds": us,
                        }
                        for node_id, (calls, us) in frame.samples.items()
                    },
                    "added_nodes": frame.added_nodes,
                    "added_budget_groups": frame.added_budget_groups,
                }
            )
        out["frames"] = frames_out

    out["summary"] = build_summary(recording, top_n)
    return out


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(
        description="Parse a Source Engine .vprof recording and emit JSON.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("input", help="path to the .vprof file")
    parser.add_argument(
        "-o", "--output", help="output file (default: stdout)", default=None
    )
    parser.add_argument(
        "--indent",
        type=int,
        default=None,
        help="JSON indentation; omit for compact output",
    )
    parser.add_argument(
        "--no-frames",
        action="store_true",
        help="omit the per-frame timing samples",
    )
    parser.add_argument(
        "--no-nodes",
        action="store_true",
        help="omit the node definition table",
    )
    parser.add_argument(
        "--summary-only",
        action="store_true",
        help="emit only metadata, budget groups, and the summary",
    )
    parser.add_argument(
        "--top",
        type=int,
        default=25,
        help="number of hottest nodes to include in the summary",
    )
    parser.add_argument(
        "--max-frames",
        type=int,
        default=None,
        help="stop after this many frames (useful for huge captures)",
    )

    args = parser.parse_args(argv)

    try:
        with open(args.input, "rb") as fh:
            raw = fh.read()
    except OSError as e:
        sys.stderr.write(f"error: cannot read {args.input}: {e}\n")
        return 2

    try:
        recording = parse_vprof(raw)
    except VProfError as e:
        sys.stderr.write(f"error: {e}\n")
        return 1

    if args.max_frames is not None and args.max_frames >= 0:
        recording.frames = recording.frames[: args.max_frames]

    include_frames = not (args.no_frames or args.summary_only)
    include_nodes = not args.no_nodes

    result = recording_to_dict(
        recording,
        source_name=args.input,
        size_bytes=len(raw),
        include_frames=include_frames,
        include_nodes=include_nodes,
        top_n=args.top,
    )

    text = json.dumps(result, indent=args.indent, ensure_ascii=False)
    if args.output:
        with open(args.output, "w", encoding="utf-8") as fh:
            fh.write(text)
    else:
        try:
            sys.stdout.write(text)
            sys.stdout.write("\n")
            sys.stdout.flush()
        except BrokenPipeError:
            # e.g. `vprof_to_json.py file.vprof | head`
            try:
                sys.stdout.close()
            except OSError:
                pass
            return 0

    return 0


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