import re

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

extra = r'''
__G3N = 0
function __G3(k, v)
  if __G3N < 3000 then
    __G3N = __G3N + 1
    local vt = type(v)
    local vs
    if vt == "string" then vs = string.format("%q", v:sub(1, 60))
    elseif vt == "number" or vt == "boolean" or vt == "nil" then vs = tostring(v)
    elseif vt == "table" then
      local ok, n = pcall(function() local z = 0 for _ in pairs(v) do z = z + 1 end return z end)
      vs = "table#" .. tostring(ok and n or "?")
    else vs = vt end
    __realprint("G3", tostring(k), vt, vs)
  end
end
__SETN = 0
function __G4(name, v)
  if __SETN < 1000 then
    __SETN = __SETN + 1
    local vt = type(v)
    local vs = vt == "string" and string.format("%q", v:sub(1,60)) or (vt == "table" and ("table#" .. tostring(#v)) or tostring(v))
    __realprint("SETG", tostring(name), vt, vs)
  end
end
__PLOGN = 0
function __PLOG(name, k)
  if __PLOGN < 1000 then
    __PLOGN = __PLOGN + 1
    __realprint("PROXY", tostring(name), tostring(k))
  end
end
__OPN = 0
__OPDISTINCT = 0
__OPSEEN = {}
__OPDONE = false
function __OP(m, q, z)
  __OPN = __OPN + 1
  local s = __OPSEEN[m]
  if not s then
    s = {}
    __OPSEEN[m] = s
    local ins = -1
    if type(m) == "table" and type(m[3]) == "table" then ins = #m[3] end
    __realprint("PROTO", tostring(m), "ins", ins, "m2", tostring(m[2]))
  end
  if not s[q] then
    s[q] = true
    __OPDISTINCT = __OPDISTINCT + 1
    if __OPDISTINCT <= 4000 then
      local ins = -1
      if type(m) == "table" and type(m[3]) == "table" then ins = #m[3] end
      __realprint("OP", tostring(m), q, z, "ins", ins)
    end
  end
end
'''

old = '''    __index = function(self, k)
      if k == "ClassName" then return "Proxy" end'''
new = '''    __index = function(self, k)
      if __PLOG then __PLOG(name, k) end
      if k == "ClassName" then return "Proxy" end'''
shim = shim.replace(old, new, 1)
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=...;'):]

get = '(q)[d[Q]]=S[p[Q]];'
assert s.count(get) == 1
s = s.replace(get, '  local __sv=S[p[Q]];__G3(p[Q],__sv);(q)[d[Q]]=__sv;', 1)

setg = 'i=m[P[Q]];i[1][i[3]]=(q[t[Q]]);'
assert s.count(setg) == 1
s = s.replace(setg, 'i=m[P[Q]];__G4(tostring(i[3]),q[t[Q]]);i[1][i[3]]=(q[t[Q]]);', 1)

# SETFIELD global (opcode 72 / 79 / 58)
f1 = '(m[P[Q]])[p[Q]]=q[d[Q]];'
print('f1', s.count(f1))
s = s.replace(f1, ' __G4(tostring(p[Q]),q[d[Q]]);(m[P[Q]])[p[Q]]=q[d[Q]];')
f2 = '(m[P[Q]])[p[Q]]=y[Q];'
print('f2', s.count(f2))
s = s.replace(f2, ' __G4(tostring(p[Q]),y[Q]);(m[P[Q]])[p[Q]]=y[Q];')
f3 = 'm[P[Q]][q[t[Q]]]=q[d[Q]];'
print('f3', s.count(f3))
s = s.replace(f3, ' __G4(tostring(q[t[Q]]),q[d[Q]]);m[P[Q]][q[t[Q]]]=q[d[Q]];')

# opcode trace at loop head
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(m,Q,Z);', 1)

s = s.replace('return({', 'local __T={', 1)
s = s.rstrip()[:-len('}):A();')] + '};__WRAP(__T);return __T:A();'

profile = r'''
local __CNT = {}
local __CALLS = 0
function __WRAP(T)
  for k, v in pairs(T) do
    if type(v) == "function" then
      __CNT[k] = 0
      T[k] = function(...)
        __CNT[k] = __CNT[k] + 1
        __CALLS = __CALLS + 1
        return v(...)
      end
    end
  end
end
'''

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))
  local ks = {}
  for k, n in pairs(__CNT) do ks[#ks+1] = k end
  table.sort(ks, function(a, b) return __CNT[a] > __CNT[b] end)
  __realprint("CALLS", __CALLS, "OPS", __OPN, "DISTINCT", __OPDISTINCT)
  for i = 1, math.min(#ks, 20) do __realprint("HOT", ks[i], __CNT[ks[i]]) end
end
'''
out = shim + profile + tail
open('/tmp/opencode/probe4.luau', 'w').write(out)
print("wrote", len(out))
