import re, sys, os

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

extra = r'''
-- extra probes
__DYN_N = 0
function __DYN(k, c)
  if __DYN_N < 400 then
    __DYN_N = __DYN_N + 1
    local t = type(c)
    local n = -1
    if t == "table" then
      local ok, v = pcall(function() return #c end)
      if ok then n = v end
    end
    __realprint("DYN", tostring(k), t, n)
  end
end
__GL_N = 0
function __GLOG(kind, k, cell, val)
  if __GL_N < 400 then
    __GL_N = __GL_N + 1
    local ok, name = pcall(function() return tostring(cell[3]) end)
    local vs
    local vt = type(val)
    if vt == "string" then vs = string.format("%q", val:sub(1, 40))
    elseif vt == "number" or vt == "boolean" then vs = tostring(val)
    elseif vt == "table" then
      local ok2, n = pcall(function() local z = 0 for _ in pairs(val) do z = z + 1 end return z end)
      vs = "table#" .. tostring(ok2 and n or "?")
    else vs = vt end
    __realprint("GLOB", kind, tostring(k), ok and name or "?", vs)
  end
end
__PLOGN = 0
function __PLOG(name, k)
  if __PLOGN < 300 then
    __PLOGN = __PLOGN + 1
    __realprint("PROXY", tostring(name), tostring(k))
  end
end
'''

# add proxy logging to the __proxy __index
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'''
assert old in shim
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)
esc = ''.join('\\x%02x' % b for b in rbytes)
rlit = '"' + esc + '"'

s = raw
assert s.startswith('local R=')
s = 'local R=buffer.fromstring(' + rlit + ');' + s[len('local R=...;'):]

# patch dynamic helper call site
call = 'i=W[p[Q]](c);(l[7])(i,S);q[P[Q]]=i;'
assert s.count(call) == 1
s = s.replace(call, '__DYN(p[Q],c);' + call, 1)

# patch GETGLOBAL opcode 118
get = 'i=(m[P[Q]]);q[d[Q]]=(i[1][i[3]]);'
assert s.count(get) == 1
s = s.replace(get, 'i=(m[P[Q]]);__GLOG("GET",P[Q],i,i[1][i[3]]);q[d[Q]]=(i[1][i[3]]);', 1)

# patch SETGLOBAL opcode 148
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]];__GLOG("SET",P[Q],i,q[t[Q]]);i[1][i[3]]=(q[t[Q]]);', 1)

# patch global-field get opcode 223
gf = 'i=m[P[Q]];(q)[d[Q]]=i[1][i[3]][p[Q]];'
assert s.count(gf) == 1
s = s.replace(gf, 'i=m[P[Q]];__GLOG("GETFIELD",P[Q],i,i[1][i[3]][p[Q]]);(q)[d[Q]]=i[1][i[3]][p[Q]];', 1)

# wrap the table for method counting, cap high
s = s.replace('return({', 'local __T={', 1)
s = s.rstrip()[:-len('}):A();')] + '};__WRAP(__T);return __T:A();'

profile = r'''
local __CNT = {}
local __FIRST = {}
local __LAST = {}
local __CALLS = 0
local __FN2NAME = {}
local __realerror = error
function __WRAP(T)
  for k, v in pairs(T) do
    if type(v) == "function" then __FN2NAME[v] = tostring(k) end
  end
  for k, v in pairs(T) do
    if type(v) == "function" then
      __CNT[k] = 0
      __FIRST[k] = {}
      T[k] = function(...)
        __CNT[k] = __CNT[k] + 1
        __CALLS = __CALLS + 1
        if __CALLS % 2000000 == 0 then __realprint("BEAT", __CALLS, k) end
        local n = #__FIRST[k]
        if n < 4 then
          local parts = {}
          local a = table.pack(...)
          for i = 1, math.min(4, a.n) do
            local x = a[i]
            parts[#parts+1] = type(x) == "number" and string.format("%.8g", x) or type(x) == "string" and string.format("%q", x:sub(1, 20)) or type(x)
          end
          __FIRST[k][n+1] = table.concat(parts, ",")
        end
        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)
  for i = 1, math.min(#ks, 20) do
    local k = ks[i]
    __realprint("HOT", k, __CNT[k], table.concat(__FIRST[k], " || "))
  end
end
'''

out = shim + profile + tail
open('/tmp/opencode/probe2.luau', 'w').write(out)
print("wrote", len(out))