#!/usr/bin/env python3
"""Find the timers responsible for the size of a Paje (.trace) file.

Paje format recap (see the header of any *.trace file):

  5  Alias Type Name Color         -> PajeDefineEntityValue: maps a value alias
                                       (e.g. a30) to a human-readable timer Name.
  6  Time Alias Type Container Name -> PajeCreateContainer:  maps a container alias
                                       (e.g. a14) to a Name ("Thread 0", ...).
  12 Time Type Container Value Id   -> PajePushState: opens a timer. `Value` is the
                                       value alias -> the timer Name we care about.
  13 Time Type Container            -> PajePopState: closes the innermost open timer
                                       on that container (no value of its own).

The file is completely dominated by 12/13 lines, so a timer's contribution to the
file size is (bytes of all its push lines) + (bytes of the matching pop lines).
Pops don't name a value, so we attribute each pop to the timer at the top of a
per-container stack.

Usage:  python3 filter.py [trace-file]   (default: 2_default.trace)
"""

import sys
from collections import defaultdict


def human(n):
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024 or unit == "TB":
            return f"{n:.1f} {unit}" if unit != "B" else f"{n} B"
        n /= 1024


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else "2_default.trace"

    value_name = {}       # value alias  -> timer Name          (from "5" lines)
    container_name = {}   # container alias -> container Name    (from "6" lines)

    # per-timer accumulators, keyed by value alias
    n_push = defaultdict(int)
    bytes_total = defaultdict(int)   # push + attributed pop bytes

    # per-(container, value) accumulators
    n_push_cv = defaultdict(int)
    bytes_cv = defaultdict(int)

    # per-container stack of currently-open value aliases
    stacks = defaultdict(list)

    unmatched_pops = 0

    with open(path, "rb") as f:
        for raw in f:
            n = len(raw)                       # exact bytes this line occupies
            code = raw[:2]

            if code == b"12":                  # PajePushState
                # 12 \t Time \t Type \t Container \t Value \t Id
                parts = raw.split(b"\t")
                container = parts[3]
                value = parts[4]
                n_push[value] += 1
                bytes_total[value] += n
                n_push_cv[container, value] += 1
                bytes_cv[container, value] += n
                stacks[container].append(value)

            elif code == b"13":                # PajePopState
                # 13 \t Time \t Type \t Container
                parts = raw.split(b"\t")
                container = parts[3].rstrip(b"\r\n")
                st = stacks[container]
                if st:
                    value = st.pop()
                    bytes_total[value] += n
                    bytes_cv[container, value] += n
                else:
                    unmatched_pops += 1

            elif code[:1] == b"5" and raw[1:2] in (b" ", b"\t"):
                # 5 Alias Type Name Color   (Name is a "..."-quoted string)
                parts = raw.split(b"\t")
                if len(parts) >= 4:
                    alias = parts[1].strip()
                    name = parts[3].strip().strip(b'"')
                    value_name[alias] = name.decode("utf-8", "replace")

            elif code[:1] == b"6" and raw[1:2] in (b" ", b"\t"):
                # 6 Time Alias Type Container Name
                parts = raw.split(b"\t")
                if len(parts) >= 6:
                    alias = parts[2].strip()
                    name = parts[5].strip().strip(b'"')
                    container_name[alias] = name.decode("utf-8", "replace")

    # ---- report -----------------------------------------------------------
    grand_pushes = sum(n_push.values())
    grand_bytes = sum(bytes_total.values())

    rows = sorted(bytes_total.items(), key=lambda kv: kv[1], reverse=True)

    print(f"File: {path}")
    print(f"Total push events: {grand_pushes:,}")
    print(f"Bytes in push+pop lines: {human(grand_bytes)} "
          f"({grand_bytes:,} bytes)")
    if unmatched_pops:
        print(f"Unmatched pops (ignored): {unmatched_pops:,}")
    print()

    TOP = 25
    print(f"Top {TOP} space-consuming timers:")
    print(f"{'bytes':>12} {'share':>7} {'pushes':>14}  timer name")
    print("-" * 100)
    for alias, nbytes in rows[:TOP]:
        name = value_name.get(alias, alias.decode())
        share = 100.0 * nbytes / grand_bytes if grand_bytes else 0.0
        print(f"{human(nbytes):>12} {share:6.2f}% {n_push[alias]:>14,}  {name}")

    # ---- per-container summary --------------------------------------------
    per_container = defaultdict(int)
    for (container, _value), nbytes in bytes_cv.items():
        per_container[container] += nbytes

    print()
    print("Bytes by container:")
    print(f"{'bytes':>12} {'share':>7}  container")
    print("-" * 60)
    for container, nbytes in sorted(per_container.items(),
                                    key=lambda kv: kv[1], reverse=True):
        cname = container_name.get(container, container.decode())
        share = 100.0 * nbytes / grand_bytes if grand_bytes else 0.0
        print(f"{human(nbytes):>12} {share:6.2f}%  {cname}")

    # ---- top (container, timer) pairs -------------------------------------
    print()
    print(f"Top {TOP} (container, timer) pairs:")
    print(f"{'bytes':>12} {'share':>7} {'pushes':>14}  container: timer name")
    print("-" * 100)
    pairs = sorted(bytes_cv.items(), key=lambda kv: kv[1], reverse=True)
    for (container, value), nbytes in pairs[:TOP]:
        cname = container_name.get(container, container.decode())
        name = value_name.get(value, value.decode())
        share = 100.0 * nbytes / grand_bytes if grand_bytes else 0.0
        print(f"{human(nbytes):>12} {share:6.2f}% "
              f"{n_push_cv[container, value]:>14,}  {cname}: {name}")


if __name__ == "__main__":
    main()