#!/usr/bin/env python3
from pathlib import Path
import subprocess

ROOT = Path("/tmp/opencode")
loader = (Path("/opt/ws-remote/luraph_pack/loader_guard.luau")).read_text()
start = loader.index('local B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZ')
end = loader.index("local function sample(s)")
codec = loader[start:end]

prelude = r'''
local byte, char, concat = string.byte, string.char, table.concat
local band, bxor, rshift, lshift = bit32.band, bit32.bxor, bit32.rshift, bit32.lshift
local wait = function() end
'''

tail = r'''
local B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
local function b64enc(s)
    local out, cur, cn = {}, {}, 0
    local i, len = 1, #s
    while i + 2 <= len do
        local a, b, c = byte(s, i, i + 2)
        local x = a * 65536 + b * 256 + c
        cn = cn + 1
        cur[cn] = string.sub(B64, band(rshift(x, 18), 63) + 1, band(rshift(x, 18), 63) + 1)
            .. string.sub(B64, band(rshift(x, 12), 63) + 1, band(rshift(x, 12), 63) + 1)
            .. string.sub(B64, band(rshift(x, 6), 63) + 1, band(rshift(x, 6), 63) + 1)
            .. string.sub(B64, band(x, 63) + 1, band(x, 63) + 1)
        if cn >= 4096 then out[#out + 1] = concat(cur); cur = {}; cn = 0 end
        i = i + 3
    end
    local rem = len - i + 1
    if rem > 0 then
        local a, b = byte(s, i, i + 1)
        local x = a * 65536 + (b or 0) * 256
        cn = cn + 1
        cur[cn] = string.sub(B64, band(rshift(x, 18), 63) + 1, band(rshift(x, 18), 63) + 1)
            .. string.sub(B64, band(rshift(x, 12), 63) + 1, band(rshift(x, 12), 63) + 1)
            .. (rem == 2 and string.sub(B64, band(rshift(x, 6), 63) + 1, band(rshift(x, 6), 63) + 1) or "=")
            .. "="
    end
    out[#out + 1] = concat(cur)
    return concat(out)
end

local n = 619155
local t = {}
for i = 1, n do t[i] = char((i * 37) % 256) end
local data = concat(t)
local key = "0123456789abcdef0123456789abcdef"
local enc = b64enc(data)
print("sizes", #data, #enc)

local t0 = os.clock()
local decoded = b64(enc, #data)
print("b64", string.format("%.4f", os.clock() - t0), decoded and #decoded, decoded == data)

local t1 = os.clock()
local out = decrypt(data, key)
print("dec", string.format("%.4f", os.clock() - t1), out and #out)
assert(out and #out == #data, "decrypt length")
for p = 1, #data, 997 do
    local ref = bxor(byte(data, p), byte(key, ((p - 1) % 32) + 1), p % 256, math.floor(p / 256) % 256)
    if byte(out, p) ~= ref then error("mismatch at " .. p) end
end
print("ok")
'''

for label, extra in (("buffer", ""), ("fallback", "local buffer = nil\n")):
    src = prelude + extra + codec + tail
    path = ROOT / f"test_loader_codec_{label}.luau"
    path.write_text(src, encoding="utf-8")
    proc = subprocess.run([str(ROOT / "luau-src/build/luau"), str(path)],
                          capture_output=True, text=True, timeout=120)
    print("===", label)
    print(proc.stdout.strip())
    if proc.stderr.strip():
        print("STDERR:", proc.stderr.strip()[:2000])