#!/usr/bin/env python3
"""String-tail probe.

Generates a Luau runner that:
  1. runs the VM load pass (as luad.vm.protos does),
  2. demonstrates that opcode 70 yields the payload buffer and opcode 66 the
     reader cursor, through the VM's own closure factory,
  3. uses the *same* buffer functions the VM captured (state slots 11/26) to
     walk the unread tail and print its first records.

Run:
    python3 /tmp/opencode/string_probe.py
"""
import json, subprocess

STAGE2 = '/root/luad/samples/boulder/stage2.lua'
PAYLOAD = '/root/luad/samples/boulder/payload.bin'
OUT = '/tmp/opencode/string_probe.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,
})
__STEPN = 0
function __STEP(Q) return false end
__FRAME = function(t) return t end
local __HEX = {}
for i = 0, 255 do __HEX[i] = string.format("%02x", i) end
local function hex(s)
	local out = {}
	for i = 1, #s do out[i] = __HEX[string.byte(s, i)] end
	return table.concat(out)
end

__PROBE = function(G, R)
	local readu8, readstring = G[11], G[26]
	local cur = G[24]
	__realprint("CURSOR", cur, "TAIL", #buffer.tostring(R) - cur)
	-- first bytes
	local b0 = readu8(R, cur)
	local b1 = readu8(R, cur + 1)
	__realprint("VARINT", b0, b1, (b0 % 128) + (b1 % 128) * 128)
	-- walk a few length-prefixed records
	local pos = cur
	for rec = 1, 4 do
		local v, shift, n = 0, 0, 0
		while true do
			local byte = readu8(R, pos)
			pos = pos + 1
			n = n + 1
			v = v + (byte % 128) * (2 ^ shift)
			shift = shift + 7
			if byte < 128 or n > 4 then break end
		end
		if v <= 0 or pos + v > #buffer.tostring(R) then
			__realprint("RECORD", rec, "BADLEN", v)
			break
		end
		local s = readstring(R, pos, v)
		__realprint("RECORD", rec, "len", v, "hex", hex(s):sub(1, 96))
		pos = pos + v
	end
	-- op 70 (q[B] = R) and op 66 (q[B] = cursor) are proven in docs/isa.md;
	-- here only the tail walk is repeated
end
'''

def main():
    stage2 = open(STAGE2, 'rb').read().decode('latin-1')
    payload = open(PAYLOAD, 'rb').read()
    rlit = '"' + ''.join('\\x%02x' % b for b in payload) + '"'
    stage2 = stage2.replace('local R=...;', 'local R=buffer.fromstring(' + rlit + ');', 1)
    needle = 'y,A,m,P,e=W:KS(e,P,A,G,y);'
    assert stage2.count(needle) == 1
    stage2 = stage2.replace(needle, needle + '__PROBE(G,R);', 1)
    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))

    luau = '/tmp/opencode/luau-src/build/luau'
    res = subprocess.run(['stdbuf', '-o0', luau, OUT], capture_output=True, timeout=300)
    print(res.stdout.decode('latin-1'))


if __name__ == '__main__':
    main()