#!/usr/bin/env python3
"""Native goalie-only full-slapshot audit; all saves/states/results remain private.

Example: verify_shot_balance.py current.bin /private/audit --seed preferences.srm
For matched builds, pass --fixtures-from the preceding audit's output directory.
This resets positions, player control, transient carry state and PRNG before the
release; it never changes native player ratings, shot probability or goalie AI.
See tools/notes/SHOT_BALANCE_AUDIT.md for the limits of goal-rate comparisons.
"""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import argparse, hashlib, json, os, struct, subprocess

HERE = Path(__file__).resolve().parent
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument('rom', type=Path)
ap.add_argument('output', type=Path)
ap.add_argument('--seed', type=Path, required=True)
ap.add_argument('--fixtures-from', type=Path)
ap.add_argument('--teams', type=int, nargs='+', default=[10, 6, 26])
ap.add_argument('--difficulties', type=int, nargs='+', default=[2, 0, 1])
ap.add_argument('--periods', type=int, nargs='+', default=[0, 1, 2, 3])
ap.add_argument('--shots', type=int, default=200)
ap.add_argument('--jobs', type=int, default=3)
a = ap.parse_args()
assert a.shots >= 10 and a.shots % 2 == 0
assert set(a.teams) <= set(range(32))
assert set(a.difficulties) <= {0, 1, 2} and set(a.periods) <= {0, 1, 2, 3}
out = a.output.resolve()
out.mkdir(parents=True, exist_ok=True)
rom = a.rom.resolve()
# Freeze the exact hooks used by every subprocess in this run.
for filename in ['shot_balance_hook.py', 'shot_balance_fixture_hook.py']:
    destination = out / filename
    blob = (HERE / filename).read_bytes()
    if destination.exists():
        assert destination.read_bytes() == blob, 'use a fresh output folder after a hook change'
    destination.write_bytes(blob)


def run_core(folder, hook, variables, save=None):
    env = {k: v for k, v in os.environ.items() if not k.startswith('GPGX_')}
    env.update(GPGX_DEV='513', GPGX_HOOK=str(hook), **variables)
    if save:
        env['GPGX_SRM'] = str(save)
    with (folder / 'log.txt').open('w') as log:
        subprocess.run(['python3', str(HERE / 'test_core.py'), str(rom),
                        str(folder / 'commands.txt'), str(folder)], env=env,
                       stdout=log, stderr=log, check=True, timeout=180)


def fixture(team, difficulty, period):
    spec = dict(team=team, difficulty=difficulty, period=period)
    folder = (a.fixtures_from.resolve() if a.fixtures_from else out) / f'fixture-{team}-{difficulty}-{period}'
    if a.fixtures_from:
        assert (folder / 'base.st').is_file()
        assert json.loads((folder / 'fixture.json').read_text())['spec'] == spec
        return folder
    folder.mkdir(exist_ok=True)
    raw = a.seed.read_bytes()
    assert len(raw) == 65536, 'seed must be an explicit 64KB RetroArch save copy'
    logical = bytearray(raw[1::2])
    logical[0xe20:0x1ef6] = bytes(0x1ef6 - 0xe20)
    struct.pack_into('>4s6H', logical, 0x1000, b'SESN', 0, 0, 0, 84, 1, 3)
    logical[0x1fff] = sum(logical[:0x1ffe]) % 256
    logical[0x1ffe] = logical[0x1fff] ^ 255
    private = bytearray(65536)
    private[1::2] = logical
    save = folder / 'private.srm'
    save.write_bytes(private)
    (folder / 'commands.txt').write_text('18000 exit\n')
    run_core(folder, out / 'shot_balance_fixture_hook.py', {'FIXTURE': json.dumps(spec)}, save)
    assert (folder / 'base.st').is_file(), (folder, 'native game did not reach live play')
    return folder


def shoot(folder, spec):
    folder.mkdir(exist_ok=True)
    (folder / 'shots.jsonl').unlink(missing_ok=True)
    (folder / 'commands.txt').write_text('9999999 exit\n')
    run_core(folder, out / 'shot_balance_hook.py', {'SHOT_CASE': json.dumps(spec)})
    result = json.loads((folder / 'results.json').read_text())
    shots = result['shots']
    assert len(shots) == spec['count']
    assert len({s['rng_release_output'] for s in shots}) >= len(shots) * .95
    assert all(s['charge'] == 31 and s['animation'] == '0x7fc' and s['shooter_slot'] == spec['slot'] for s in shots)
    assert all((s['distance_actual'] >= 160) == (s['distance'] >= 180) for s in shots)
    assert all(not s['bonus'] for s in shots if s['distance_actual'] < 160)
    if result['fixture']['difficulty'] == 2:
        assert all(s['crit'] for s in shots)
    return result


def case(values):
    team, difficulty, period = values
    source = fixture(team, difficulty, period)
    meta = json.loads((source / 'fixture.json').read_text())
    slot = next(actor['slot'] for actor in meta['actors'][:5] if actor['role'] == 2)
    spec = dict(fixture=str(source), method=5, slot=slot, count=a.shots,
                charge_at=100, distances=[120, 200])
    result = shoot(out / f'case-{team}-{difficulty}-{period}', spec)
    assert meta['home'] == team and meta['difficulty'] == difficulty and meta['period'] == period
    print('PASS', values, 'goals', sum(s['result'] == 'goal' for s in result['shots']), flush=True)
    return result


cases = [(t, d, p) for t in a.teams for d in a.difficulties for p in a.periods]
with ThreadPoolExecutor(max_workers=a.jobs) as pool:
    results = list(pool.map(case, cases))
report = dict(rom_sha256=hashlib.sha256(rom.read_bytes()).hexdigest(),
              hook_sha256=hashlib.sha256((out / 'shot_balance_hook.py').read_bytes()).hexdigest(),
              cases=results)
(out / 'results.json').write_text(json.dumps(report, indent=2) + '\n')
print('PASS', len(cases), 'native cases /', sum(len(r['shots']) for r in results), 'full forehand slapshots')