#!/usr/bin/env python3
"""Headless critical-rating/shot regression; all generated ROMs stay in work_dir.

Usage: verify_critical.py current.bin work_dir [--baseline preceding.bin]
Requires the cached SGDK assembler and the same Genesis Plus GX core as test_core.py.
Uses isolated reset-vector drivers: no emulator window and no user save access.
The shot fixture constructs the native C2F2 stack frame, then executes the real
C4A8 power/velocity tail, including native register restoration. Lookup tests
exercise all 832 identities through both the shared lookup and attribute hook.
"""
from pathlib import Path
import argparse
import json
import os
import struct
import subprocess
from critical import TABLE, ratings

here = Path(__file__).resolve().parent


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('rom', type=Path)
    ap.add_argument('work_dir', type=Path)
    ap.add_argument('--baseline', type=Path)
    args = ap.parse_args()
    work = args.work_dir.resolve()
    work.mkdir(parents=True, exist_ok=True)
    current = args.rom.read_bytes()
    assert current[TABLE:TABLE+832] == ratings(current)[0], 'generated rating table differs'

    def run(name, source, rom, edits=()):
        out = work/name
        out.mkdir(exist_ok=True)
        asm = out/'driver.s'
        asm.write_text(source)
        obj = out/'driver-code.bin'
        subprocess.run([str(here.parent/'m68kasm.sh'), str(asm), '2FEA00', str(obj)], check=True, stdout=subprocess.DEVNULL)
        code = obj.read_bytes()
        data = bytearray(rom)
        assert len(code) < 0x600 and all(x == 255 for x in data[0x2FEA00:0x2FEA00+len(code)])
        data[0x2FEA00:0x2FEA00+len(code)] = code
        struct.pack_into('>I', data, 4, 0x2FEA00)
        for address, blob in edits:
            data[address:address+len(blob)] = blob
        test_rom = out/'driver.bin'
        test_rom.write_bytes(data)
        cmd = out/'commands.txt'
        cmd.write_text('1000 dump fff600 16 result\n1000 dump ff0000 832 lookup\n1000 dump ffd066 4 random\n1001 exit\n')
        env = {k:v for k,v in os.environ.items() if k not in ['GPGX_SRM', 'GPGX_HOOK', 'GPGX_WAV']}
        env['GPGX_DEV'] = '513'
        subprocess.run(['nice', '-n', '19', 'python3', str(here/'test_core.py'), str(test_rom), str(cmd), str(out)], env=env, check=True, stdout=subprocess.DEVNULL)
        results = {}
        for line in cmd.with_suffix('.out').read_text().splitlines():
            _, key, hexbytes = line.split(' ', 2)
            results[key] = bytes.fromhex(hexbytes)
        assert results['result'][8:12] == b'CRIT', (name, 'driver did not finish')
        return results

    source = (here/'testcases/critical_driver.s').read_text()
    cases = [('zero', 0, 0, 13), ('half', 50, 0, 13), ('max', 100, 0, 13),
             ('away-half', 50, 64, 12), ('home-reversed', 100, 128, 13),
             ('away-reversed', 50, 192, 12), ('other-skater', 0, 0, 12),
             ('goalie', 0, 0, 0), ('invalid-slot', 0, 0, 26)]
    report = []
    for name, stat, flags, slot in cases:
        s = source.replace('move.b #13,0x66', f'move.b #{slot},0x66').replace('move.b #0,0x62', f'move.b #{flags},0x62')
        if flags & 128:
            s = s.replace('#-200', '#200').replace('#-16320', '#16320').replace('#-13056', '#13056')
        # Two known identities have different ratings, to catch use of attack
        # direction instead of team membership or of on-ice index instead of slot.
        edits = [(TABLE+26*10+13, bytes([stat if not flags & 64 else 100])),
                 (TABLE+26*25+12, bytes([stat if flags & 64 else 50]))]
        if name == 'other-skater':
            edits += [(TABLE+26*10+12, b'\0'), (TABLE+26*10+13, b'\x64')]
        result = run(name, s, current, edits)
        ordinary, critical, errors = struct.unpack_from('>3H', result['result'])
        assert errors == 0 and ordinary+critical == 10000, (name, ordinary, critical, errors)
        if not stat:
            assert critical == 0 and result['random'] == bytes.fromhex('12345678'), (name, 'zero chance consumed RNG or boosted')
        else:
            assert abs(critical/10000-((stat+30) if stat else 0)/1000) < .015, (name, critical)   # r26: +3 points for every eligible shot
        report.append(dict(case=name, rating=stat, shots=10000, critical=critical, native_velocity_checks='passed'))
    if args.baseline:
        result = run('baseline', source.replace('jmp 0x2FF400', 'jmp 0xC4A8'), args.baseline.read_bytes())
        assert struct.unpack_from('>3H', result['result']) == (10000, 0, 0)
        report.append(dict(case='baseline', shots=10000, critical=0, native_velocity_checks='passed'))
    result = run('all-lookups', (here/'testcases/critical_lookup.s').read_text(), current, [(0x8E4C, bytes.fromhex('4e75'))])
    assert result['lookup'] == current[TABLE:TABLE+832]
    assert result['result'][6:8] == b'\0\0', 'lookup/attribute/register mismatch'
    report.append(dict(case='all-lookups', identities=832, native_attribute_path='passed', register_preservation='passed'))
    (work/'results.json').write_text(json.dumps(report, indent=2)+'\n')
    print('PASS', len(report), 'cases: per-player/side/direction probability, native power/velocity and all attribute lookups')


if __name__ == '__main__':
    main()