#!/usr/bin/env python3
"""Disassemble the AMD "F32" blobs (RLC, SDMA) that this tree could never read.

Background: `psp_out/FIRMWARE-COVERAGE-MATRIX.md` listed RLC and SDMA as **blocked** --
fixed-width 32-bit words, NOP 0x90000000, and no public disassembler (`f32_probe.py`
verified that rather than inheriting the claim). fail0verflow's `radeon-tools`
contains one, written for Sea Islands (CI) CP/RLC. Upstream says SDMA is "a related
but incompatible variant"; on gfx11 it decodes fine, so that caveat does not hold here.

This wrapper does the three things the upstream script cannot:

1. **Container parsing.** Upstream takes a raw image and hardcodes "skip 0x100 for
   lowercase filenames". A gfx11 `rlc` container holds *eight* separately PSP-loadable
   sub-blobs (RLC_G, IRAM, DRAM, RLCP, RLCV, and the save/restore lists), each with
   its own nested `$PS1` header, and SDMA v2_0 holds TH0 + TH1.
2. **Load bases.** A blob whose code is not loaded at word 0 disassembles with every
   call target out of range. RLC_G's 915 targets land in 0x2004..0x37ba -- its own
   size shifted by exactly 0x2000 -- so its base is 0x2000. Recovered automatically.
3. **A decode-quality metric with controls**, because the decoder is permissive: raw
   "percentage decoded" gives 78% on an RV64 image, so it proves nothing on its own.
   What separates real F32 from noise is `ret` count and whether call targets resolve
   to in-image labels.

Requires `./fetch_ref.sh` first (upstream is not vendored: no LICENSE in that repo).

usage:
  ./f32_audit.py --selftest                     # controls: CI blobs vs RV64 vs random
  ./f32_audit.py info psp_out/gc_11_0_0_rlc.bin
  ./f32_audit.py dis psp_out/sdma_6_0_0.bin --blob TH0 | less
"""
import argparse
import os
import re
import struct
import subprocess
import sys
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
TOOL = os.path.join(HERE, 'ref', 'radeon-tools', 'f32', 'f32dis.py')


def need_tool():
    if not os.path.exists(TOOL):
        raise SystemExit('missing %s -- run %s/fetch_ref.sh first' % (TOOL, HERE))


def _ps1(d, off):
    """Every PSP-loadable payload carries a nested 0x100-byte $PS1 header."""
    return d[off + 0x10:off + 0x14] == b'$PS1'


def slice_blob(d, off, size):
    if not size or off + size > len(d):
        return None
    return d[off + (0x100 if _ps1(d, off) else 0):off + size]


def split_rlc(d):
    """gfx11 rlc_firmware_header_v2_3 -> the separately loadable sub-blobs."""
    ucsz, ucoff = struct.unpack_from('<II', d, 0x14)
    p = 0x20 + 4 * 18                      # past common header + v2_0
    v21 = struct.unpack_from('<13I', d, p)
    p += 4 * 13
    v22 = struct.unpack_from('<4I', d, p)
    p += 4 * 4
    v23 = struct.unpack_from('<8I', d, p)
    return [('RLC_G', ucoff, ucsz), ('IRAM', v22[1], v22[0]), ('DRAM', v22[3], v22[2]),
            ('RLCP', v23[3], v23[2]), ('RLCV', v23[7], v23[6]),
            ('SRL_CNTL', v21[4], v21[3]), ('SRL_GPM', v21[8], v21[7]),
            ('SRL_SRM', v21[12], v21[11])]


def split_sdma(d):
    """sdma_firmware_header_v2_0: ctx (TH0) and ctl (TH1) halves."""
    ucsz, ucoff = struct.unpack_from('<II', d, 0x14)
    try:
        ctx_ucode_size, ctx_jt_off, ctx_jt_size, ctl_ucode_off, ctl_ucode_size = \
            struct.unpack_from('<5I', d, 0x24)
        if ctl_ucode_off and ctl_ucode_off + ctl_ucode_size <= len(d):
            return [('TH0', ucoff, ctx_ucode_size or ucsz),
                    ('TH1', ctl_ucode_off, ctl_ucode_size)]
    except struct.error:
        pass
    return [('TH0', ucoff, ucsz)]


def split(path):
    d = open(path, 'rb').read()
    base = os.path.basename(path).lower()
    if 'rlc' in base:
        return d, split_rlc(d)
    if 'sdma' in base:
        return d, split_sdma(d)
    ucsz, ucoff = struct.unpack_from('<II', d, 0x14)
    return d, [('UCODE', ucoff, ucsz)]


def disassemble(body):
    """Run upstream f32dis on a payload; returns its stdout."""
    need_tool()
    with tempfile.TemporaryDirectory() as td:
        # lowercase name => upstream skips 0x100 bytes and decodes little-endian
        p = os.path.join(td, 'blob')
        open(p, 'wb').write(bytes(0x100) + body)
        return subprocess.run([sys.executable, TOOL, p], capture_output=True,
                              text=True, timeout=1800).stdout


def metrics(out):
    lines = [l.strip() for l in out.splitlines() if l.startswith('    ')]
    labels = {l[:-1] for l in out.splitlines() if l.endswith(':')}
    n = len(lines)
    unknown = sum(1 for l in lines if 'dw 0x' in l)
    rets = sum(1 for l in lines if l == 'ret')
    calls = [l for l in lines if re.match(r'^(?:bl|b)\s', l)]
    raw = [int(m.group(1), 16) for l in calls
           for m in [re.match(r'^(?:bl|b)\s+0x([0-9a-f]+)', l)] if m]
    named = sum(1 for l in calls if l.split()[1].strip() in labels)
    base = 0
    if raw:
        cands = [0] + [1 << k for k in range(8, 21)] + [(min(raw) // 0x1000) * 0x1000]
        base = max(cands, key=lambda b: sum(1 for t in raw if 0 <= t - b < n))
    in_range = named + (sum(1 for t in raw if 0 <= t - base < n) if raw else 0)
    return {'insns': n, 'unknown': unknown, 'ret': rets, 'calls': len(calls),
            'resolved': in_range, 'base': base,
            'decoded_pct': 100.0 * (n - unknown) / max(1, n),
            'resolved_pct': 100.0 * in_range / max(1, len(calls))}


def verdict(m):
    """F32 code, data, or not-this-ISA -- deliberately conservative."""
    if m['insns'] == 0:
        return 'empty'
    if m['ret'] == 0:
        # No return instruction anywhere => not a callable code image. Data blobs
        # still yield stray "calls" because random words decode as branches, so
        # ret is the discriminator, not call count. (IRAM/DRAM/SRL_* land here,
        # which matches their names: save/restore register lists, not code.)
        return 'data (no returns; %.0f%% undecodable)' % (100 - m['decoded_pct'])
    if m['resolved_pct'] >= 95 and m['ret'] > 0:
        return 'F32 CODE (base 0x%x)' % m['base']
    if m['resolved_pct'] >= 80:
        return 'probably F32, some targets unresolved'
    return 'not F32 / undetermined'


def cmd_info(path):
    d, blobs = split(path)
    print('%s (%d bytes)' % (path, len(d)))
    print('  %-9s %-9s %-9s %-6s %7s %7s %7s %9s  %s' % (
        'blob', 'offset', 'size', '$PS1', 'insns', 'unk', 'ret', 'calls-ok', 'verdict'))
    for name, off, size in blobs:
        body = slice_blob(d, off, size)
        if body is None:
            print('  %-9s 0x%-7x 0x%-7x  <absent>' % (name, off, size))
            continue
        m = metrics(disassemble(body))
        print('  %-9s 0x%-7x 0x%-7x %-6s %7d %7d %7d %8.1f%%  %s' % (
            name, off, size, _ps1(d, off), m['insns'], m['unknown'], m['ret'],
            m['resolved_pct'], verdict(m)))


def cmd_dis(path, want):
    d, blobs = split(path)
    for name, off, size in blobs:
        if want and name != want:
            continue
        body = slice_blob(d, off, size)
        if body is None:
            continue
        out = disassemble(body)
        m = metrics(out)
        print('### %s %s  base=0x%x  %d insns  %s' % (
            os.path.basename(path), name, m['base'], m['insns'], verdict(m)))
        print(out)


def selftest():
    ok = True

    def chk(c, what):
        nonlocal ok
        ok = ok and bool(c)
        print('  [%s] %s' % ('ok' if c else 'FAIL', what))

    need_tool()
    refs = os.path.join(HERE, 'ref')
    pos = os.path.join(refs, 'BONAIRE_rlc.bin')
    if not os.path.exists(pos):
        print('  (skipping positive control: %s not fetched)' % pos)
    else:
        # uppercase name: upstream reads it big-endian from offset 0, as intended
        out = subprocess.run([sys.executable, TOOL, pos], capture_output=True,
                             text=True, timeout=600).stdout
        m = metrics(out)
        chk(m['decoded_pct'] > 99, 'positive control BONAIRE_rlc decodes %.1f%% > 99%%'
            % m['decoded_pct'])
        chk(m['ret'] > 10, 'positive control has %d ret instructions' % m['ret'])

    import random
    random.seed(1)
    m = metrics(disassemble(bytes(random.randrange(256) for _ in range(0x4000))))
    chk(m['decoded_pct'] < 40, 'negative control (random) decodes only %.1f%%'
        % m['decoded_pct'])
    chk(m['ret'] == 0, 'negative control (random) has no ret instructions')
    chk('not F32' in verdict(m) or 'data' in verdict(m),
        'negative control verdict is not "F32 CODE" (got: %s)' % verdict(m))

    print('\n  note: raw decode%% is NOT sufficient evidence -- an RV64 image scores')
    print('  78%% here. ret count and call-target resolution are what discriminate.')
    print('\nSELFTEST %s' % ('PASS' if ok else 'FAIL'))
    return 0 if ok else 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('cmd', nargs='?', choices=['info', 'dis'])
    ap.add_argument('image', nargs='?')
    ap.add_argument('--blob', default=None)
    ap.add_argument('--selftest', action='store_true')
    a = ap.parse_args()
    if a.selftest:
        raise SystemExit(selftest())
    if not a.cmd or not a.image:
        ap.error('need a command and an image (or --selftest)')
    if a.cmd == 'info':
        cmd_info(a.image)
    else:
        cmd_dis(a.image, a.blob)
    return 0


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