#!/usr/bin/env python3
"""Generate a Luau runner that single-steps synthetic one-instruction
prototypes through the VM's own closure factory and records every effect.

Patches (all on a copy in /tmp/opencode; samples untouched):
  * payload embedded as a \\xNN literal
  * `while true do local Z=o[Q];` -> step hook (halt before 2nd instruction)
  * `local q,M=(l[41](A));`      -> frame proxy for read/write tracing
  * after the load pass: __HARNESS(G,R)
"""
import json
import sys

STAGE2 = '/root/luad/samples/boulder/stage2.lua'
PAYLOAD = '/root/luad/samples/boulder/payload.bin'
DUMP = '/tmp/opencode/boulder-protos.json'
OUT = '/tmp/opencode/exec_run.luau'

SHIM = r'''
local __realprint = print
tick = function() return 0 end
time = function() return 0 end
warn = function() end
typeof = function(v) return type(v) end
local __HOST = getfenv()
__ENV = setmetatable({}, {
	__index = function(t, k) return __HOST[k] end,
	__newindex = function(t, k, v) __HOST[k] = v end,
})

-- ---------------------------------------------------------------- tracing --
__ON = false
__BUDGET = 1
__STEPN = 0
__EV = {}
__TEST = 0

function __LOG(...)
	if __ON then __EV[#__EV + 1] = table.pack(...) end
end

function __STEP(Q)
	if not __ON then return false end
	__STEPN = __STEPN + 1
	__LOG("PC", Q)
	if __STEPN > __BUDGET then __LOG("HALT"); return true end
	return false
end

local __utag = setmetatable({}, { __mode = "k" })
local function __V(v)
	local t = type(v)
	if t == "table" then return "P:" .. tostring(__utag[v] or "tbl") end
	if t == "string" then return "s:" .. v:sub(1, 24) end
	if t == "number" then return "n:" .. string.format("%.10g", v) end
	if t == "boolean" then return "b:" .. tostring(v) end
	if t == "nil" then return "nil" end
	return t
end
__VV = __V

local __UMT
__U = function(tag)
	local p = setmetatable({}, __UMT)
	__utag[p] = tag
	return p
end
__UMT = {
	__index = function(t, k) __LOG("IDX", tostring(__utag[t]), tostring(k)); return __U(tostring(k)) end,
	__newindex = function(t, k, v) __LOG("SIDX", tostring(__utag[t]), tostring(k), __V(v)) end,
	__call = function(t, ...) __LOG("CALL", tostring(__utag[t]), select("#", ...)); return __U("ret") end,
	__add = function(a, b) __LOG("OP", "add"); return __U("+") end,
	__sub = function(a, b) __LOG("OP", "sub"); return __U("-") end,
	__mul = function(a, b) __LOG("OP", "mul"); return __U("*") end,
	__div = function(a, b) __LOG("OP", "div"); return __U("/") end,
	__idiv = function(a, b) __LOG("OP", "idiv"); return __U("//") end,
	__mod = function(a, b) __LOG("OP", "mod"); return __U("%") end,
	__pow = function(a, b) __LOG("OP", "pow"); return __U("^") end,
	__unm = function(a) __LOG("OP", "unm"); return __U("-x") end,
	__concat = function(a, b) __LOG("OP", "concat"); return __U("..") end,
	__len = function(t) __LOG("OP", "len"); return 7 end,
	__eq = function(a, b) __LOG("OP", "eq"); return false end,
	__lt = function(a, b) __LOG("OP", "lt"); return false end,
	__le = function(a, b) __LOG("OP", "le"); return false end,
	__tostring = function(t) return "<" .. tostring(__utag[t]) .. ">" end,
}

__SEED = {}
__FRAME = function(t)
	if not __ON then return t end
	local store = __SEED
	return setmetatable({}, {
		__index = function(_, k) local v = store[k]; __LOG("R", k, __V(v)); return v end,
		__newindex = function(_, k, v) __LOG("W", k, __V(v)); store[k] = v end,
		__len = function() __LOG("OP", "framelen"); return 8 end,
	})
end

local __KV = {}
__KPROXY = function(slot, value)
	local p = setmetatable({}, {
		__index = function(_, k) __LOG("K" .. slot, k, __V(value)); return value end,
	})
	return p
end

__KCOUNT = { 0, 0, 0 }
__HARNESS = function(G, R)
	local host_pcall = pcall
	local host_unpack = table.unpack
	local host_select = select
	local format = string.format

	-- swallow instruction errors so the interpreter's error path never runs
	local orig_pcall = G[4]
	G[4] = function(f, ...)
		local r = table.pack(host_pcall(f, ...))
		if not r[1] then __LOG("ERR", tostring(r[2])); return true end
		return host_unpack(r, 1, r.n)
	end

	-- log every state-helper call: slot, argument count, first args
	local orig = {}
	for s = 1, 60 do
		local f = rawget(G, s)
		if type(f) == "function" and s ~= 4 and s ~= 55 then
			orig[s] = f
			G[s] = function(...)
				__LOG("S", s, host_select("#", ...))
				return f(...)
			end
		end
	end
	local factory = G[55]

	local OPS = __OPS__
	local modes = { "proxy", "num", "str", "buf" }

	local function seed_for(mode)
		local seed = {}
		if mode == "proxy" then
			for i = 1, 24 do seed[i] = __U("r" .. i) end
			return seed, __U("K7"), __U("K8"), __U("K9")
		elseif mode == "num" then
			for i = 1, 24 do seed[i] = i * 10 end
			return seed, 700, 800, 900
		elseif mode == "str" then
			for i = 1, 24 do seed[i] = "s" .. i end
			return seed, "k7", "k8", "k9"
		else -- buf: offsets/counts are real numbers, one register holds R
			for i = 1, 24 do seed[i] = i end
			seed[1] = 100
			seed[2] = 4
			seed[5] = R
			seed[6] = G[24]
			return seed, 1, 2, 3
		end
	end

	local function run(op, a, b, c, mode)
		__TEST = __TEST + 1
		local id = __TEST
		__EV = {}
		__STEPN = 0
		local seed, k7, k8, k9 = seed_for(mode)
		__SEED = seed
		local p = {
			[1] = 0, [2] = {},
			[3] = { op }, [4] = { b }, [5] = { c },
			[6] = 24,
			[7] = __KPROXY(7, k7), [8] = __KPROXY(8, k8), [9] = __KPROXY(9, k9),
			[10] = { a }, [11] = {},
		}
		local up = { { [1] = __U("U1"), [3] = 1 } }
		local cursor0 = G[24]
		__ON = true
		local ok, err = host_pcall(function()
			local fn = factory(p, up)
			local ok2, res = host_pcall(fn)
			if not ok2 then __LOG("ERR", tostring(res)) end
		end)
		__ON = false
		if not ok then __LOG("ERR", tostring(err)) end
		local function fmt(v)
			if type(v) == "number" then
				if v == math.floor(v) and math.abs(v) < 1e12 then return format("%d", v) end
				return format("%.10g", v)
			elseif type(v) == "table" then return __V(v)
			elseif type(v) == "boolean" then return tostring(v)
			else return tostring(v) end
		end
		__realprint("B", id, op, mode, cursor0)
		for i = 1, #__EV do
			local e = __EV[i]
			local parts = { "E", id }
			for j = 1, e.n do parts[#parts + 1] = fmt(e[j]) end
			__realprint(table.concat(parts, "\t"))
		end
		__realprint("E", id, "CURSOR", cursor0, G[24])
	end

	for i = 1, #OPS do
		local op = OPS[i]
		for m = 1, #modes do
			run(op, 1, 2, 3, modes[m])
		end
	end
	__realprint("HARNESS_END", __TEST)
	G[4] = orig_pcall
	for s, f in pairs(orig) do G[s] = f end
end
'''

def main():
    stage2 = open(STAGE2, 'rb').read().decode('latin-1')
    payload = open(PAYLOAD, 'rb').read()
    dump = json.load(open(DUMP))
    ops = sorted({ins['op'] for p in dump['protos'] for ins in p['code']})
    rlit = '"' + ''.join('\\x%02x' % b for b in payload) + '"'
    stage2 = stage2.replace('local R=...;', 'local R=buffer.fromstring(' + rlit + ');', 1)
    for needle, repl, cnt in [
        ('while true do local Z=o[Q];', 'while true do if __STEP(Q)then return end;local Z=o[Q];', 1),
        ('local q,M=(l[41](A));', 'local q,M=__FRAME((l[41](A)));', 1),
    ]:
        assert stage2.count(needle) == cnt, (needle, stage2.count(needle))
        stage2 = stage2.replace(needle, repl, 1)
    needle = 'y,A,m,P,e=W:KS(e,P,A,G,y);'
    assert stage2.count(needle) == 1
    stage2 = stage2.replace(needle, needle + '__HARNESS(G,R);', 1)
    ops_lua = '{' + ','.join(str(o) for o in ops) + '}'
    shim = SHIM.replace('__OPS__', ops_lua)
    out = shim + 'local __ok,__err=pcall(function()\nlocal __fn=function()\n' + stage2 + '\nend\nsetfenv(__fn,__ENV)\nreturn __fn()\nend)\n'
    out += '__realprint("EXEC", __ok, tostring(__err))\n'
    open(OUT, 'w').write(out)
    print("wrote", OUT, len(out), "opcodes", len(ops))

if __name__ == '__main__':
    main()