#!/usr/bin/env python3
"""Build a split-handler VM runtime with lazily decrypted handler closures.

The original VM's dispatch chain is parsed into individual handler bodies.
Each body is wrapped in a function factory, encrypted, and stored in a handler
table.  A handler is decrypted and compiled only on its first use.
"""
from __future__ import annotations

import random
import re


def _lua_quote(text: str) -> str:
    data = text.encode("utf-8", "surrogateescape")
    out = ['"']
    for b in data:
        if b == 92:
            out.append("\\\\")
        elif b == 34:
            out.append('\\"')
        elif 10 <= b < 13 or b == 9:
            out.append("\\%03d" % b)
        elif 32 <= b < 127:
            out.append(chr(b))
        else:
            out.append("\\%03d" % b)
    out.append('"')
    return "".join(out)


def _stream(data: bytes, seed: int) -> bytes:
    state = (seed ^ 0x9E3779B9) & 0xFFFFFFFF
    out = bytearray(data)
    for i in range(len(out)):
        state = (state ^ ((state << 13) & 0xFFFFFFFF)) & 0xFFFFFFFF
        state = (state ^ (state >> 17)) & 0xFFFFFFFF
        state = (state ^ ((state << 5) & 0xFFFFFFFF)) & 0xFFFFFFFF
        out[i] ^= state & 0xFF
    return bytes(out)


def _encrypt_source(source: str, seed: int) -> str:
    return _stream(source.encode("utf-8", "surrogateescape"), seed).hex()


def _extract_handlers(vm: str):
    lines = vm.splitlines(True)
    start = None
    for i, line in enumerate(lines):
        if line.startswith('\t\tif op =='):
            start = i
            break
    if start is None:
        raise ValueError("handler chain not found")
    else_at = None
    for i in range(start + 1, len(lines)):
        if lines[i].startswith('\t\telse\n') or lines[i].startswith('\t\telse '):
            else_at = i
            break
    if else_at is None:
        raise ValueError("handler chain end not found")
    starts = [i for i in range(start, else_at)
              if lines[i].startswith('\t\tif op ==') or lines[i].startswith('\t\telseif op ==')]
    handlers = []
    for n, at in enumerate(starts):
        end = starts[n + 1] if n + 1 < len(starts) else else_at
        header = lines[at]
        body = "".join(lines[at + 1:end])
        ops = [int(v, 0) for v in re.findall(r'op\s*==\s*(0x[0-9A-Fa-f]+|\d+)', header)]
        if not ops:
            continue
        terminal = bool(re.search(r'(?m)^\s*return\b', body))
        handlers.append((ops, body, terminal))
    return handlers


def build_handler_runtime(vm: str, key: int = 0) -> str:
    """Return a runtime with one lazily compiled closure per opcode."""
    handlers = _extract_handlers(vm)
    rng = random.Random((int(key) * 2654435761 + 0x51ED270B) & 0xFFFFFFFF)
    records = []
    for ops, body, terminal in handlers:
        seed = rng.randrange(0x10000, 0x7FFFFFFF)
        source = (
            "return function(ctx, ins)\n"
            "local q,V,top,ups,va,nva,rb=ctx.q,ctx.V,ctx.top,ctx.ups,ctx.va,ctx.nva,ctx.rb\n"
            "local K,pack,unpack,cell,cellval,cellput,mkclos,truthy,ENV,callwith="
            "ctx.K,ctx.pack,ctx.unpack,ctx.cell,ctx.cellval,ctx.cellput,ctx.mkclos,ctx.truthy,ctx.ENV,ctx.callwith\n"
            "local pc=ctx.pc\n"
            "if ctx.q==nil then return end\n"
            + body +
            "\nctx.pc=pc\nctx.top=top\nreturn\nend\n"
        ) if not terminal else (
            "return function(ctx, ins)\n"
            "local q,V,top,ups,va,nva,rb=ctx.q,ctx.V,ctx.top,ctx.ups,ctx.va,ctx.nva,ctx.rb\n"
            "local K,pack,unpack,cell,cellval,cellput,mkclos,truthy,ENV,callwith="
            "ctx.K,ctx.pack,ctx.unpack,ctx.cell,ctx.cellval,ctx.cellput,ctx.mkclos,ctx.truthy,ctx.ENV,ctx.callwith\n"
            "local pc=ctx.pc\n"
            "if ctx.q==nil then return end\n"
            + body + "\nend\n"
        )
        encrypted = _encrypt_source(source, seed)
        for op in ops:
            records.append((op, terminal, seed, encrypted))
    rng.shuffle(records)
    table = "{\n" + ",\n".join(
        "[%d]={t=%s,s=%d,x=%s}" % (op, "true" if terminal else "false", seed, _lua_quote(enc))
        for op, terminal, seed, enc in records
    ) + "\n}"
    # Prefix: everything through callwith/mkclos, including binary decoder.
    exec_at = vm.find("function V.exec(")
    if exec_at < 0:
        raise ValueError("V.exec not found")
    suffix_at = vm.find("local M = {}")
    if suffix_at < 0:
        raise ValueError("M not found")
    decode_at = vm.find("-- Binary graph reader.")
    prefix = vm[:exec_at]
    suffix = vm[decode_at:] if decode_at >= 0 else vm[suffix_at:]
    new_exec = r'''
local H = %s

local function h_unhex(s)
	local out = {}
	for i = 1, #s, 2 do
		out[#out + 1] = string.char(tonumber(s:sub(i, i + 1), 16))
	end
	return table.concat(out)
end

local function h_decrypt(hex, seed)
	local s = h_unhex(hex)
	local state = bit32.bxor(seed, 0x9E3779B9)
	local out = {}
	for i = 1, #s do
		state = bit32.bxor(state, bit32.lshift(state, 13))
		state = bit32.bxor(state, bit32.rshift(state, 17))
		state = bit32.bxor(state, bit32.lshift(state, 5))
		out[i] = string.char(bit32.bxor(s:byte(i), bit32.band(state, 255)))
	end
	return table.concat(out)
end

local function h_get(op)
	local rec = H[op]
	if not rec then error('', 0) end
	if not rec.fn then
		local src = h_decrypt(rec.x, rec.s)
		local f = loadstring or load
		local factory = f(src)
		if not factory then error('', 0) end
		rec.fn = factory()
	end
	return rec
end

function V.exec(pid, ups, ...)
	local p = V.protos[pid]
	if not p then error('', 0) end
	local q = {}
	local rb = V.regbase or 0
	local np = p.params
	local nargs = select('#', ...)
	for i = 1, np do
		q[rb + i - 1] = (select(i, ...))
	end
	local va, nva = {}, 0
	if p.vararg then
		nva = nargs - np
		for i = 1, nva do
			va[i] = select(np + i, ...)
		end
	end
	local code = p.code
	local n = #code
	local pc = 1
	local top = rb - 1
	local ctx = {
		q = q, V = V, ups = ups, va = va, nva = nva, rb = rb,
		top = top, pc = pc, K = K, pack = pack, unpack = unpack,
		cell = cell, cellval = cellval, cellput = cellput,
		mkclos = mkclos, truthy = truthy, ENV = ENV, callwith = callwith,
	}
	while pc >= 1 and pc <= n do
		local ins = code[pc]
		local op = V.opmap and V.opmap[ins[1]] or bit32.bxor(ins[1], V.opkey)
		local rec = h_get(op)
		ctx.pc, ctx.top = pc, top
		local pr = pack(pcall(rec.fn, ctx, ins))
		if not pr[1] then error('handler ' .. tostring(op) .. ': ' .. tostring(pr[2]), 0) end
		local r = pack(unpack(pr, 2, pr.n))
		if rec.t then
			return unpack(r, 1, r.n)
		end
		pc, top = ctx.pc, ctx.top
	end
end
''' % table
    return prefix + new_exec + suffix