#!/usr/bin/env python3
"""
KGSL GPU Snapshot Parser
Parses kgsl_snapshot_* binary files produced by the Adreno/KGSL driver.
"""

import sys
import os
import struct
from ctypes import *
from collections import namedtuple

# ─── Constants ───────────────────────────────────────────────────────────────

SNAPSHOT_MAGIC          = 0x504D0002
SNAPSHOT_SECTION_MAGIC  = 0xABCD

SECTION_NAMES = {
    0x0101: "OS",
    0x0201: "REGS",
    0x0202: "REGS_V2",
    0x0302: "RB_V2",
    0x0402: "IB_V2",
    0x0501: "INDEXED_REGS",
    0x0502: "INDEXED_REGS_V2",
    0x0901: "DEBUG",
    0x0A01: "DEBUGBUS",
    0x0B02: "GPU_OBJECT_V2",
    0x0E02: "MEMLIST_V2",
    0x1201: "SHADER",
    0x1202: "SHADER_V2",
    0x1203: "SHADER_V3",
    0x1501: "MVC",
    0x1502: "MVC_V2",
    0x1503: "MVC_V3",
    0x1601: "GMU",
    0x1701: "GMU_MEMORY",
    0x1801: "SIDE_DEBUGBUS",
    0x1901: "TRACE_BUFFER",
    0x1A01: "EVENTLOG",
    0xFFFF: "END",
}

GMU_MEM_TYPES = {
    0x00: "UNKNOWN",
    0x01: "HFI",
    0x02: "LOG",
    0x03: "BWTABLE",
    0x04: "DEBUG",
    0x05: "BIN_BLOCK",
    0x06: "CONTEXT_QUEUE",
    0x07: "HW_FENCE",
    0x08: "WARMBOOT",
    0x09: "VRB",
    0x0a: "TRACE",
}

GPU_OBJECT_TYPES = {
    1: "SHADER",
    2: "IB",
    3: "GENERIC",
    4: "DRAW",
    5: "GLOBAL",
}

DEBUG_TYPES = {
    7:  "CP_MEQ",
    8:  "CP_PM4_RAM",
    9:  "CP_PFP_RAM",
    10: "CP_ROQ",
    11: "SHADER_MEMORY",
    12: "CP_MERCIU",
    14: "SQE_VERSION",
    15: "GMU_CORE_VERSION",
    16: "GMU_CORE_DEV_VERSION",
    17: "GMU_PWR_VERSION",
    18: "GMU_PWR_DEV_VERSION",
    19: "GMU_HFI_VERSION",
    20: "AQE_VERSION",
    32: "SLICE_MASK",
}

# ─── Structures (must match kernel packing) ──────────────────────────────────

class SnapshotHeader(Structure):
    _pack_ = 1
    _fields_ = [
        ("magic",  c_uint32),
        ("gpuid",  c_uint32),
        ("chipid", c_uint32),
    ]

class SectionHeader(Structure):
    _pack_ = 1
    _fields_ = [
        ("magic", c_uint16),
        ("id",    c_uint16),
        ("size",  c_uint32),
    ]

class RbV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("start",             c_int32),
        ("end",               c_int32),
        ("rbsize",            c_int32),
        ("wptr",              c_int32),
        ("rptr",              c_int32),
        ("count",             c_int32),
        ("timestamp_queued",  c_uint32),
        ("timestamp_retired", c_uint32),
        ("gpuaddr",           c_uint64),
        ("id",                c_uint32),
    ]

class GmuMem(Structure):
    _pack_ = 1
    _fields_ = [
        ("type",     c_int32),
        ("hostaddr", c_uint64),
        ("gmuaddr",  c_uint64),
        ("gpuaddr",  c_uint64),
    ]

class IbV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("gpuaddr", c_uint64),
        ("ptbase",  c_uint64),
        ("size",    c_uint64),
    ]

class MemListV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("num_entries", c_int32),
        ("ptbase",      c_uint64),
    ]
    
class Regs(Structure):
    _pack_ = 1
    _fields_ = [
        ("count", c_uint32),
    ]    

class Gpu(Structure):
    _pack_ = 1
    _fields_ = [
        ("count", c_uint32),
    ]

class IndexedRegs(Structure):
    _pack_ = 1
    _fields_ = [
        ("index_reg", c_uint32),
        ("data_reg",  c_uint32),
        ("start",     c_int32),
        ("count",     c_int32),
    ]

class IndexedRegsV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("index_reg", c_uint32),
        ("data_reg",  c_uint32),
        ("start",     c_uint32),
        ("count",     c_uint32),
        ("pipe_id",   c_uint32),
        ("slice_id",  c_uint32),
    ]

class GpuObjectV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("type",    c_int32),
        ("gpuaddr", c_uint64),
        ("ptbase",  c_uint64),
        ("size",    c_uint64),   # in dwords
    ]

class Debug(Structure):
    _pack_ = 1
    _fields_ = [
        ("type", c_int32),
        ("size", c_int32),   # in dwords
    ]

class Debugbus(Structure):
    _pack_ = 1
    _fields_ = [
        ("id",    c_int32),
        ("count", c_int32),
    ]

class Shader(Structure):
    _pack_ = 1
    _fields_ = [
        ("type",  c_int32),
        ("index", c_int32),
        ("size",  c_int32),
    ]

class ShaderV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("type",     c_int32),
        ("index",    c_int32),
        ("usptp",    c_int32),
        ("pipe_id",  c_int32),
        ("location", c_int32),
        ("size",     c_uint32),
    ]

class LinuxOsV4(Structure):
    _pack_ = 1
    _fields_ = [
        ("osid",                    c_int32),
        ("seconds",                 c_uint32),
        ("power_flags",             c_uint32),
        ("power_level",             c_uint32),
        ("power_interval_timeout",  c_uint32),
        ("grpclk",                  c_uint32),
        ("busclk",                  c_uint32),
        ("ptbase",                  c_uint64),
        ("ptbase_lpac",             c_uint64),
        ("pid",                     c_uint32),
        ("pid_lpac",                c_uint32),
        ("current_context",         c_uint32),
        ("current_context_lpac",    c_uint32),
        ("ctxtcount",               c_uint32),
        ("release",                 c_char * 32),
        ("version",                 c_char * 32),
        ("comm",                    c_char * 16),
        ("comm_lpac",               c_char * 16),
    ]

class LinuxContextV2(Structure):
    _pack_ = 1
    _fields_ = [
        ("id",                  c_uint32),
        ("timestamp_queued",    c_uint32),
        ("timestamp_consumed",  c_uint32),
        ("timestamp_retired",   c_uint32),
    ]

# ─── Helpers ─────────────────────────────────────────────────────────────────

def u32_list(data):
    """Interpret raw bytes as a list of little-endian uint32."""
    n = len(data) // 4
    return list(struct.unpack(f"<{n}I", data[:n*4]))

def hexdump(data, width=8, max_dwords=64):
    """Print a compact dword hexdump, truncated if large."""
    dwords = u32_list(data)
    if not dwords:
        print("    (empty)")
        return
    truncated = len(dwords) > max_dwords
    if truncated:
        dwords = dwords[:max_dwords]
    for i in range(0, len(dwords), width):
        chunk = dwords[i:i+width]
        hexstr = " ".join(f"{x:08x}" for x in chunk)
        print(f"    {i*4:06x}: {hexstr}")
    if truncated:
        print(f"    ... ({len(u32_list(data))} dwords total, showing first {max_dwords})")

def safe_str(b):
    try:
        return b.split(b"\x00")[0].decode("ascii", errors="replace")
    except Exception:
        return repr(b)

# ─── Section parsers ─────────────────────────────────────────────────────────

def parse_os(payload):
    if len(payload) < sizeof(LinuxOsV4):
        print("  [OS] payload too small")
        return
    hdr = LinuxOsV4.from_buffer_copy(payload)
    print(f"  OS ID          : 0x{hdr.osid:08x}")
    print(f"  Timestamp      : {hdr.seconds}")
    print(f"  Power flags    : 0x{hdr.power_flags:08x}")
    print(f"  Power level    : {hdr.power_level}")
    print(f"  GRPCLK         : {hdr.grpclk}")
    print(f"  BUSCLK         : {hdr.busclk}")
    print(f"  ptbase         : 0x{hdr.ptbase:016x}")
    print(f"  ptbase_lpac    : 0x{hdr.ptbase_lpac:016x}")
    print(f"  PID            : {hdr.pid}  ({safe_str(hdr.comm)})")
    print(f"  PID_LPAC       : {hdr.pid_lpac}  ({safe_str(hdr.comm_lpac)})")
    print(f"  Current ctxt   : {hdr.current_context}")
    print(f"  Current ctxt LPAC: {hdr.current_context_lpac}")
    print(f"  Kernel release : {safe_str(hdr.release)}")
    print(f"  Kernel version : {safe_str(hdr.version)}")
    print(f"  Context count  : {hdr.ctxtcount}")

    offset = sizeof(LinuxOsV4)
    for i in range(hdr.ctxtcount):
        if offset + sizeof(LinuxContextV2) > len(payload):
            break
        ctx = LinuxContextV2.from_buffer_copy(payload[offset:])
        print(f"    ctx[{i}] id={ctx.id}  queued={ctx.timestamp_queued}  "
              f"consumed={ctx.timestamp_consumed}  retired={ctx.timestamp_retired}")
        offset += sizeof(LinuxContextV2)

def parse_rb_v2(payload):
    if len(payload) < sizeof(RbV2):
        print("  [RB] payload too small")
        return
    rb = RbV2.from_buffer_copy(payload)
    print(f"  RB id          : {rb.id}")
    print(f"  gpuaddr        : 0x{rb.gpuaddr:016x}")
    print(f"  rbsize (dwords): {rb.rbsize}")
    print(f"  start/end      : {rb.start} / {rb.end}")
    print(f"  wptr / rptr    : {rb.wptr} / {rb.rptr}")
    print(f"  count          : {rb.count}")
    print(f"  ts queued      : {rb.timestamp_queued}")
    print(f"  ts retired     : {rb.timestamp_retired}")
    data = payload[sizeof(RbV2):]
    print(f"  Ringbuffer data ({len(data)} bytes):")
    hexdump(data, max_dwords=128)

def parse_gmu_memory(payload):
    if len(payload) < sizeof(GmuMem):
        print("  [GMU_MEM] payload too small")
        return
    m = GmuMem.from_buffer_copy(payload)
    tname = GMU_MEM_TYPES.get(m.type, f"UNKNOWN({m.type})")
    print(f"  Type           : {tname} ({m.type})")
    print(f"  hostaddr       : 0x{m.hostaddr:016x}")
    print(f"  gmuaddr        : 0x{m.gmuaddr:016x}")
    print(f"  gpuaddr        : 0x{m.gpuaddr:016x}")
    data = payload[sizeof(GmuMem):]
    print(f"  Data ({len(data)} bytes):")
    hexdump(data, max_dwords=64)

def parse_ib_v2(payload):
    if len(payload) < sizeof(IbV2):
        print("  [IB] payload too small")
        return
    ib = IbV2.from_buffer_copy(payload)
    print(f"  gpuaddr        : 0x{ib.gpuaddr:016x}")
    print(f"  ptbase         : 0x{ib.ptbase:016x}")
    print(f"  size (bytes?)  : {ib.size}")
    data = payload[sizeof(IbV2):]
    print(f"  IB contents ({len(data)} bytes):")
    hexdump(data, max_dwords=64)

def parse_memlist_v2(payload):
    if len(payload) < sizeof(MemListV2):
        print("  [MEMLIST] payload too small")
        return
    ml = MemListV2.from_buffer_copy(payload)
    print(f"  num_entries    : {ml.num_entries}")
    print(f"  ptbase         : 0x{ml.ptbase:016x}")
    # Remaining data is usually pairs of (gpuaddr, size) or similar – dump raw
    data = payload[sizeof(MemListV2):]
    print(f"  Entries data ({len(data)} bytes):")
    hexdump(data, max_dwords=64)

def parse_regs(payload):
    if len(payload) < sizeof(Regs):
        print("  [REGS] payload too small")
        return
    r = Regs.from_buffer_copy(payload)
    print(f"  count (pairs)  : {r.count}")
    data = payload[sizeof(Regs):]
    # Usually stored as offset/value pairs
    dwords = u32_list(data)
    pairs = min(r.count, len(dwords)//2, 128)
    print(f"  Register dump (showing up to {pairs} pairs):")
    for i in range(pairs):
        off = dwords[i*2]
        val = dwords[i*2+1]
        print(f"    [{off:08x}] = {val:08x}")
    if r.count > pairs:
        print(f"    ... ({r.count} pairs total)")

def parse_indexed_regs(payload, v2=False):
    if v2:
        if len(payload) < sizeof(IndexedRegsV2):
            print("  [INDEXED_REGS_V2] too small")
            return
        ir = IndexedRegsV2.from_buffer_copy(payload)
        print(f"  index_reg      : 0x{ir.index_reg:08x}")
        print(f"  data_reg       : 0x{ir.data_reg:08x}")
        print(f"  start          : {ir.start}")
        print(f"  count          : {ir.count}")
        print(f"  pipe_id        : {ir.pipe_id}")
        print(f"  slice_id       : {ir.slice_id}")
        data = payload[sizeof(IndexedRegsV2):]
    else:
        if len(payload) < sizeof(IndexedRegs):
            print("  [INDEXED_REGS] too small")
            return
        ir = IndexedRegs.from_buffer_copy(payload)
        print(f"  index_reg      : 0x{ir.index_reg:08x}")
        print(f"  data_reg       : 0x{ir.data_reg:08x}")
        print(f"  start          : {ir.start}")
        print(f"  count          : {ir.count}")
        data = payload[sizeof(IndexedRegs):]
    print(f"  Data ({len(data)} bytes):")
    hexdump(data, max_dwords=64)

def parse_gpu_object_v2(payload):
    if len(payload) < sizeof(GpuObjectV2):
        print("  [GPU_OBJECT] too small")
        return
    obj = GpuObjectV2.from_buffer_copy(payload)
    tname = GPU_OBJECT_TYPES.get(obj.type, f"UNKNOWN({obj.type})")
    print(f"  type           : {tname}")
    print(f"  gpuaddr        : 0x{obj.gpuaddr:016x}")
    print(f"  ptbase         : 0x{obj.ptbase:016x}")
    print(f"  size (dwords)  : {obj.size}")
    data = payload[sizeof(GpuObjectV2):]
    print(f"  Object data ({len(data)} bytes):")
    hexdump(data, max_dwords=64)

def parse_debug(payload):
    if len(payload) < sizeof(Debug):
        print("  [DEBUG] too small")
        return
    d = Debug.from_buffer_copy(payload)
    tname = DEBUG_TYPES.get(d.type, f"UNKNOWN({d.type})")
    print(f"  type           : {tname}")
    print(f"  size (dwords)  : {d.size}")
    data = payload[sizeof(Debug):]
    print(f"  Data ({len(data)} bytes):")
    hexdump(data, max_dwords=64)

def parse_debugbus(payload):
    if len(payload) < sizeof(Debugbus):
        print("  [DEBUGBUS] too small")
        return
    db = Debugbus.from_buffer_copy(payload)
    print(f"  id             : {db.id}")
    print(f"  count (dwords) : {db.count}")
    data = payload[sizeof(Debugbus):]
    print(f"  Data ({len(data)} bytes):")
    hexdump(data, max_dwords=32)

def parse_shader(payload, v2=False):
    if v2:
        if len(payload) < sizeof(ShaderV2):
            print("  [SHADER_V2] too small")
            return
        s = ShaderV2.from_buffer_copy(payload)
        print(f"  type={s.type}  index={s.index}  usptp={s.usptp}  "
              f"pipe={s.pipe_id}  loc={s.location}  size={s.size}")
        data = payload[sizeof(ShaderV2):]
    else:
        if len(payload) < sizeof(Shader):
            print("  [SHADER] too small")
            return
        s = Shader.from_buffer_copy(payload)
        print(f"  type={s.type}  index={s.index}  size={s.size}")
        data = payload[sizeof(Shader):]
    print(f"  Shader data ({len(data)} bytes):")
    hexdump(data, max_dwords=32)

# ─── Main parser ─────────────────────────────────────────────────────────────

def parse_snapshot(filename, dump_raw_dir=None):
    with open(filename, "rb") as f:
        data = f.read()

    print(f"File          : {filename}")
    print(f"Size          : {len(data)} bytes")
    print()

    if len(data) < sizeof(SnapshotHeader):
        print("File too small for snapshot header")
        return

    hdr = SnapshotHeader.from_buffer_copy(data)
    if hdr.magic != SNAPSHOT_MAGIC:
        print(f"WARNING: bad magic 0x{hdr.magic:08x} (expected 0x{SNAPSHOT_MAGIC:08x})")
        print("Continuing anyway...\n")

    core = (hdr.gpuid >> 16) & 0xFFFF
    rev  = hdr.gpuid & 0xFFFF
    print(f"=== Snapshot Header ===")
    print(f"  magic   : 0x{hdr.magic:08x}")
    print(f"  gpuid   : 0x{hdr.gpuid:08x}  (core=0x{core:04x}  rev=0x{rev:04x})")
    print(f"  chipid  : 0x{hdr.chipid:08x}")
    print()

    offset = sizeof(SnapshotHeader)
    section_num = 0

    while offset + sizeof(SectionHeader) <= len(data):
        sect = SectionHeader.from_buffer_copy(data[offset:])
        if sect.magic != SNAPSHOT_SECTION_MAGIC:
            print(f"!!! Bad section magic 0x{sect.magic:04x} at offset 0x{offset:x} – stopping")
            break

        name = SECTION_NAMES.get(sect.id, f"UNKNOWN_0x{sect.id:04x}")
        payload_size = sect.size - sizeof(SectionHeader)
        payload_start = offset + sizeof(SectionHeader)
        payload = data[payload_start : payload_start + payload_size]

        print(f"── Section {section_num}: {name} (id=0x{sect.id:04x}) ──")
        print(f"  offset in file : 0x{offset:x}")
        print(f"  total size     : {sect.size}  (payload {payload_size} bytes)")

        if sect.id == 0xFFFF:          # END
            print("  (end of snapshot)")
            break
        elif sect.id == 0x0101:        # OS
            parse_os(payload)
        elif sect.id == 0x0302:        # RB_V2
            parse_rb_v2(payload)
        elif sect.id == 0x1701:        # GMU_MEMORY
            parse_gmu_memory(payload)
        elif sect.id == 0x0402:        # IB_V2
            parse_ib_v2(payload)
        elif sect.id == 0x0E02:        # MEMLIST_V2
            parse_memlist_v2(payload)
        elif sect.id in (0x0201, 0x0202):  # REGS / REGS_V2
            parse_regs(payload)
        elif sect.id == 0x0501:        # INDEXED_REGS
            parse_indexed_regs(payload, v2=False)
        elif sect.id == 0x0502:        # INDEXED_REGS_V2
            parse_indexed_regs(payload, v2=True)
        elif sect.id == 0x0B02:        # GPU_OBJECT_V2
            parse_gpu_object_v2(payload)
        elif sect.id == 0x0901:        # DEBUG
            parse_debug(payload)
        elif sect.id == 0x0A01:        # DEBUGBUS
            parse_debugbus(payload)
        elif sect.id == 0x1201:        # SHADER
            parse_shader(payload, v2=False)
        elif sect.id in (0x1202, 0x1203):  # SHADER_V2 / V3
            parse_shader(payload, v2=True)
        else:
            print(f"  (no specialized parser – raw dump)")
            hexdump(payload, max_dwords=32)

        # Optional: dump every section payload to disk
        if dump_raw_dir:
            os.makedirs(dump_raw_dir, exist_ok=True)
            outname = f"{dump_raw_dir}/sect_{section_num:03d}_{name}_0x{sect.id:04x}.bin"
            with open(outname, "wb") as out:
                out.write(payload)
            print(f"  → wrote raw payload to {outname}")

        print()
        offset += sect.size
        section_num += 1

    print(f"Parsed {section_num} sections.")

# ─── Entry point ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <kgsl_snapshot_file> [--dump-raw <dir>]")
        sys.exit(1)

    filename = sys.argv[1]
    dump_dir = None
    if len(sys.argv) >= 4 and sys.argv[2] == "--dump-raw":
        dump_dir = sys.argv[3]

    parse_snapshot(filename, dump_raw_dir=dump_dir)