#!/usr/bin/env python3
"""READ-ONLY probe: is DCN powered, and is the DMCUB already running?

WRITES NOTHING except the PCIE_INDEX2 window register required to perform an
indirect read (the same access idiom psp_out/dce_survey.py already uses).

This is the gate for the whole Option-A plan. It answers, in one run:

  A. access path      -- can we reach DCN registers at all, and by which path
  B. DCN powered      -- do DMU registers decode (not mirror / not all-0 / all-1)
  C. DMCUB running    -- DMCUB_CNTL.DMCUB_ENABLE=1, SOFT_RESET=0, SCRATCH0.dal_fw,
                         and DMCUB_TIMER_CURRENT advancing between reads
  D. mailbox live     -- INBOX1_BASE/SIZE non-zero, RPTR/WPTR sane

Verdict decides the next step:
  * C true  -> the DMCUB is up (something else brought it up): the type-8 PoC can
               inject into the existing ring; no bring-up needed.
  * B true, C false -> DCN is powered but DMCUB is in reset: run dmcub_bringup.py.
  * B false -> DCN is power-gated or SMN-denied at our privilege level. Bring-up
               needs display power/clock work first, and no amount of register
               writing from here will help.

usage (on the rig, with tinygrad on PYTHONPATH):
    python3 dmcub_probe.py [--inst 0] [--json out.json]
"""
import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__))), '..', 'psp_out'))

import dmcub_hwregs as H  # noqa: E402

SCRATCH0_BITS = {  # union dmub_fw_boot_status, ref/dmub_cmd.h
    0: 'dal_fw', 1: 'mailbox_rdy', 2: 'optimized_init_done', 3: 'restore_required',
    4: 'defer_load', 5: 'fams_enabled', 6: 'detection_required', 7: 'hw_power_init_done',
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--inst', type=int, default=0)
    ap.add_argument('--json', default=None)
    args = ap.parse_args()

    from tinygrad import Device
    adev = Device["AMD"].iface.dev_impl
    r = H.Regs(adev=adev, inst=args.inst, dry_run=True)
    out = {'inst': args.inst}

    print('=== A. addressing')
    print('  DCE segment bases: %s' % [hex(x) for x in (r.seg or [])])
    for n in ('DMCUB_CNTL', 'DMCUB_SEC_CNTL', 'DMCUB_SCRATCH0',
              'DMCUB_INBOX1_BASE_ADDRESS', 'DMCUB_TIMER_CURRENT'):
        print('  ' + r.describe(n))

    print('  access-path validation (control + scratch write/readback):')
    path, report = r.validate()
    out['path'] = path
    out['path_report'] = {k: str(v) for k, v in report.items()}
    if path == 'none':
        print('\nVERDICT: no usable DCN register path -- the positive control or the '
              'scratch write/readback failed.\n'
              'Do not interpret all-zero DMCUB reads as "block absent": on this card '
              'they read 0 until initialised.')
        _dump(args, out)
        return 2

    print('\n=== B/C. DMCUB state')
    vals = {}
    for n in ('DMCUB_CNTL', 'DMCUB_CNTL2', 'DMCUB_SEC_CNTL', 'DMCUB_SCRATCH0',
              'DMCUB_SCRATCH1', 'DMCUB_SCRATCH7', 'DMCUB_SCRATCH15',
              'DMCUB_REGION3_CW0_BASE_ADDRESS', 'DMCUB_REGION3_CW0_TOP_ADDRESS',
              'DMCUB_REGION3_CW4_BASE_ADDRESS', 'DMCUB_REGION3_CW4_TOP_ADDRESS',
              'DMCUB_INBOX1_BASE_ADDRESS', 'DMCUB_INBOX1_SIZE',
              'DMCUB_INBOX1_RPTR', 'DMCUB_INBOX1_WPTR',
              'DCN_VM_FB_LOCATION_BASE', 'DCN_VM_FB_OFFSET'):
        try:
            vals[n] = r.read(n)
            print('  %-34s = 0x%08x' % (n, vals[n]))
        except Exception as e:                          # noqa: BLE001
            print('  %-34s = <%s>' % (n, e))
    out['regs'] = {k: v for k, v in vals.items()}

    enabled = r.get('DMCUB_CNTL', 'DMCUB_ENABLE')
    in_reset = r.get('DMCUB_CNTL2', 'DMCUB_SOFT_RESET')
    sec_reset = r.get('DMCUB_SEC_CNTL', 'DMCUB_SEC_RESET')
    s0 = vals.get('DMCUB_SCRATCH0', 0)
    bits = [n for b, n in SCRATCH0_BITS.items() if s0 >> b & 1]
    print('  DMCUB_ENABLE=%d SOFT_RESET=%d SEC_RESET=%d  boot_status bits: %s' % (
        enabled, in_reset, sec_reset, bits or 'none'))

    t = [r.read('DMCUB_TIMER_CURRENT') for _ in range(3)]
    for _ in range(2):
        time.sleep(0.01)
        t.append(r.read('DMCUB_TIMER_CURRENT'))
    advancing = len(set(t)) > 1
    print('  DMCUB_TIMER_CURRENT samples: %s -> %s' % (
        [hex(x) for x in t], 'ADVANCING' if advancing else 'static'))
    out.update(enabled=enabled, in_reset=in_reset, sec_reset=sec_reset,
               boot_bits=bits, timer_advancing=advancing)

    running = bool(enabled and not in_reset and advancing)
    mailbox = bool(vals.get('DMCUB_INBOX1_BASE_ADDRESS') and vals.get('DMCUB_INBOX1_SIZE'))
    out.update(running=running, mailbox=mailbox)

    print('\n=== VERDICT')
    if running and mailbox:
        print('  DMCUB is RUNNING with a live inbox1 ring.')
        print('  -> skip bring-up; dmcub_poc_type8.py can inject into the existing ring.')
        print('     (ring is at DMUB VA 0x%08x size 0x%x; rptr=0x%x wptr=0x%x)' % (
            vals['DMCUB_INBOX1_BASE_ADDRESS'], vals['DMCUB_INBOX1_SIZE'],
            vals['DMCUB_INBOX1_RPTR'], vals['DMCUB_INBOX1_WPTR']))
        rc = 0
    elif running:
        print('  DMCUB is running but no inbox1 ring is programmed.')
        print('  -> dmcub_bringup.py --phase mailbox (ring only).')
        rc = 1
    else:
        print('  DCN registers decode but the DMCUB is NOT running '
              '(enable=%d reset=%d timer_advancing=%s).' % (enabled, in_reset, advancing))
        print('  -> dmcub_bringup.py (full backdoor load).')
        rc = 1
    _dump(args, out)
    return rc


def _dump(args, out):
    if args.json:
        with open(args.json, 'w') as fh:
            json.dump(out, fh, indent=2, default=str)
        print('\nwrote %s' % args.json)


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