import re, sys

mode = sys.argv[1] if len(sys.argv) > 1 else "A"

SRC = open('/tmp/opencode/make_runner.py').read()
shim = re.search(r"shim = r'''(.*?)'''", SRC, re.S).group(1)

force = ''
if mode == "A":
    force = 'if n==140 and pc==8 then regs[2]=0 end'
elif mode == "B":
    force = 'if n==140 and pc==8 then regs[9]=1000000000 end'
elif mode == "C":
    force = 'if n==140 and pc==8 then regs[2]=1; regs[9]=0 end'

extra = r'''
__NINS={}
local function __g(t,k) local ok,v=pcall(function() return t[k] end); if ok then return v end return nil end
local function __nins(O) local n=__NINS[O]; if n then return n end; local o3=__g(O,3); if type(o3)=="table" then n=#o3 else n=-1 end; __NINS[O]=n; return n end
__FN=0
function __OP(O,pc,z,a,b,c,d,e,f,regs)
  local n=__nins(O)
  FORCE_BODY
end
'''.replace("FORCE_BODY", force)

shim += extra

raw = open('/tmp/opencode/stage2.lua').read()
rhex = re.search(r'__RHEX__\t([0-9a-f]+)', open('/tmp/opencode/cap4.txt').read()).group(1)
rbytes = bytes.fromhex(rhex)
rlit = '"' + ''.join('\\x%02x' % b for b in rbytes) + '"'

s = raw
s = 'local R=buffer.fromstring(' + rlit + ');' + s[len('local R=...;'):]
head = 'while true do local Z=o[Q];'
assert s.count(head) == 1
s = s.replace(head, 'while true do local Z=o[Q];__OP(O,Q,Z,t[Q],d[Q],P[Q],y[Q],b[Q],p[Q],q);', 1)

tail = r'''
local __ok,__err = pcall(function()
''' + s + r'''
end)
__realprint("EXEC1", __ok, type(__err))
if __ok then
  local __ok2,__err2 = pcall(__err)
  __realprint("EXEC2", __ok2, tostring(__err2))
end
'''
open('/tmp/opencode/force_%s.luau' % mode, 'w').write(shim + tail)
print("wrote", mode)
