#!/usr/bin/env python3
"""Minimal Python port of Go's wasm_exec.js for running crux's main.wasm in a
sandbox (wasmtime): no browser, no network, no host file access.

Usage (library):  from gojs_sandbox import GoSandbox
                  g = GoSandbox("main.wasm"); g.run(); print(g.call_global("check", "input"))
"""
import math
import os
import struct
import time

import wasmtime

UNDEF = type("Undefined", (), {"__repr__": lambda s: "undefined"})()


class JSObject(dict):
    def __getattr__(self, k):
        return self.get(k, UNDEF)


class JSArray(list):
    pass


class U8Array(bytearray):
    pass


class JSFunc:
    def __init__(self, fn, name="fn"):
        self.fn, self.name = fn, name

    def __call__(self, this, *args):
        return self.fn(this, *args)

    def __repr__(self):
        return f"<JSFunc {self.name}>"


def js_get(v, k):
    if isinstance(k, (int, float)) and isinstance(v, (list, bytearray)):
        k = int(k)
        return v[k] if 0 <= k < len(v) else UNDEF
    if k == "length" and isinstance(v, (list, bytearray, str)):
        return float(len(v))
    if isinstance(v, dict):
        return v.get(k, UNDEF)
    return UNDEF


def js_str(v):
    if v is UNDEF:
        return "undefined"
    if v is None:
        return "null"
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, float):
        return str(int(v)) if v.is_integer() else repr(v)
    return str(v)


class GoSandbox:
    def __init__(self, wasm_path, log=None):
        self.engine = wasmtime.Engine()
        self.store = wasmtime.Store(self.engine)
        self.module = wasmtime.Module.from_file(self.engine, wasm_path)
        self.log = log or (lambda *a: None)
        self.stdout = bytearray()
        self.exited = False
        self.pending = None
        self.timeouts = {}
        self.next_timeout = 1
        self.t0 = time.time()
        g = JSObject()
        self.global_ = g
        g["Object"] = JSFunc(lambda this, *a: JSObject(), "Object")
        g["Array"] = JSFunc(lambda this, *a: JSArray(), "Array")
        g["Uint8Array"] = JSFunc(lambda this, n=0, *a: U8Array(int(n)), "Uint8Array")
        g["Date"] = JSFunc(lambda this, *a: JSObject(getTimezoneOffset=JSFunc(lambda t: 0.0)), "Date")
        consts = JSObject(O_WRONLY=-1.0, O_RDWR=-1.0, O_CREAT=-1.0, O_TRUNC=-1.0,
                          O_APPEND=-1.0, O_EXCL=-1.0, O_DIRECTORY=-1.0)
        enosys = lambda this, *a: (_ for _ in ()).throw(RuntimeError("ENOSYS"))
        g["fs"] = JSObject(constants=consts, writeSync=JSFunc(self._fs_write, "writeSync"),
                           write=JSFunc(self._fs_write_cb, "write"), open=JSFunc(enosys, "open"))
        g["process"] = JSObject(pid=-1.0, ppid=-1.0, getuid=JSFunc(lambda t: -1.0),
                                getgid=JSFunc(lambda t: -1.0), geteuid=JSFunc(lambda t: -1.0),
                                getegid=JSFunc(lambda t: -1.0), getgroups=JSFunc(enosys),
                                umask=JSFunc(enosys), cwd=JSFunc(enosys), chdir=JSFunc(enosys))
        g["crypto"] = JSObject(getRandomValues=JSFunc(lambda this, a: a))
        g["performance"] = JSObject(now=JSFunc(lambda this: (time.time() - self.t0) * 1000.0))
        self.go = JSObject(_makeFuncWrapper=JSFunc(self._make_wrapper, "_makeFuncWrapper"),
                           _pendingEvent=None)
        self.values = [math.nan, 0.0, None, True, False, g, self.go]
        self.refs = [math.inf] * len(self.values)
        self.ids = {}
        self.pool = []

    # ---------------- memory helpers
    @property
    def mem(self):
        return self.memory

    def rd(self, addr, n):
        return bytes(self.memory.read(self.store, addr, addr + n))

    def wr(self, addr, data):
        self.memory.write(self.store, bytes(data), addr)

    def i64(self, addr):
        lo, hi = struct.unpack("<Ii", self.rd(addr, 8))
        return lo + hi * 4294967296

    def set_i64(self, addr, v):
        v = int(v)
        self.wr(addr, struct.pack("<II", v & 0xFFFFFFFF, (v // 4294967296) & 0xFFFFFFFF))

    def load_value(self, addr):
        f = struct.unpack("<d", self.rd(addr, 8))[0]
        if f == 0:
            return UNDEF
        if not math.isnan(f):
            return f
        return self.values[struct.unpack("<I", self.rd(addr, 4))[0]]

    def _key(self, v):
        return id(v) if isinstance(v, (dict, list, bytearray, JSFunc)) else (type(v).__name__, v)

    def store_value(self, addr, v):
        if isinstance(v, bool):
            pass
        elif isinstance(v, (int, float)) and v != 0:
            if isinstance(v, float) and math.isnan(v):
                self.wr(addr, struct.pack("<II", 0, 0x7FF80000))
                return
            self.wr(addr, struct.pack("<d", float(v)))
            return
        if v is UNDEF:
            self.wr(addr, struct.pack("<d", 0.0))
            return
        if isinstance(v, (int, float)) and not isinstance(v, bool) and v == 0:
            v = 0.0
        k = self._key(v)
        rid = self.ids.get(k)
        if rid is None:
            fixed = {("float", 0.0): 1, ("NoneType", None): 2, ("bool", True): 3, ("bool", False): 4}
            rid = fixed.get(k)
            if rid is None and v is self.global_:
                rid = 5
            if rid is None and v is self.go:
                rid = 6
        if rid is None:
            rid = self.pool.pop() if self.pool else len(self.values)
            if rid == len(self.values):
                self.values.append(v)
                self.refs.append(0)
            else:
                self.values[rid] = v
                self.refs[rid] = 0
            self.ids[k] = rid
        self.refs[rid] += 1
        flag = 0
        if isinstance(v, JSFunc):
            flag = 4
        elif isinstance(v, str):
            flag = 2
        elif isinstance(v, (dict, list, bytearray)):
            flag = 1
        self.wr(addr, struct.pack("<II", rid, 0x7FF80000 | flag))

    def load_slice(self, addr):
        return self.i64(addr), self.i64(addr + 8)

    def load_values(self, addr):
        a, n = self.load_slice(addr)
        return [self.load_value(a + i * 8) for i in range(n)]

    def load_string(self, addr):
        a, n = self.load_slice(addr)
        return self.rd(a, n).decode("utf-8", "replace")

    def sp(self):
        return self.exports["getsp"](self.store) & 0xFFFFFFFF

    # ---------------- JS host functions used by the shims
    def _fs_write(self, this, fd, buf, *a):
        self.stdout += bytes(buf)
        self.log("[go stdout]", bytes(buf).decode(errors="replace"), end="")
        return float(len(buf))

    def _fs_write_cb(self, this, fd, buf, off, length, pos, cb):
        n = self._fs_write(this, fd, buf)
        cb(None, None, n)
        return UNDEF

    def _make_wrapper(self, this, fid):
        def wrapper(this_, *args):
            event = JSObject(id=fid, this=this_, args=JSArray(args))
            self.go["_pendingEvent"] = event
            self.resume()
            return event.get("result", UNDEF)
        return JSFunc(wrapper, f"gofunc{int(fid)}")

    # ---------------- imports
    def _imports(self):
        S = self
        t_v = wasmtime.FuncType([wasmtime.ValType.i32()], [])

        def mk(fn):
            return wasmtime.Func(S.store, t_v, lambda sp: fn(sp & 0xFFFFFFFF))

        def wasmExit(sp):
            S.exited = True
            S.log("[go exit]", struct.unpack("<i", S.rd(sp + 8, 4))[0])

        def wasmWrite(sp):
            p, n = S.i64(sp + 16), struct.unpack("<i", S.rd(sp + 24, 4))[0]
            S.stdout += S.rd(p, n)

        def nanotime1(sp):
            S.set_i64(sp + 8, (time.time() - S.t0 + 1) * 1e9)

        def walltime(sp):
            ms = int(time.time() * 1000)
            S.set_i64(sp + 8, ms // 1000)
            S.wr(sp + 16, struct.pack("<i", (ms % 1000) * 1000000))

        def schedule(sp):
            tid = S.next_timeout
            S.next_timeout += 1
            S.timeouts[tid] = S.i64(sp + 8)
            S.wr(sp + 16, struct.pack("<i", tid))

        def clear(sp):
            S.timeouts.pop(struct.unpack("<i", S.rd(sp + 8, 4))[0], None)

        def rand(sp):
            a, n = S.load_slice(sp + 8)
            S.wr(a, os.urandom(n))

        def finalize(sp):
            rid = struct.unpack("<I", S.rd(sp + 8, 4))[0]
            if rid < len(S.refs) and S.refs[rid] != math.inf:
                S.refs[rid] -= 1
                if S.refs[rid] == 0:
                    v = S.values[rid]
                    S.values[rid] = None
                    S.ids.pop(S._key(v), None)
                    S.pool.append(rid)

        def stringVal(sp):
            S.store_value(sp + 24, S.load_string(sp + 8))

        def valueGet(sp):
            r = js_get(S.load_value(sp + 8), S.load_string(sp + 16))
            S.store_value(S.sp() + 32, r)

        def valueSet(sp):
            o = S.load_value(sp + 8)
            k = S.load_string(sp + 16)
            v = S.load_value(sp + 32)
            if isinstance(o, dict):
                o[k] = v
            S.log("[set]", k, repr(v)[:80])

        def valueDelete(sp):
            o = S.load_value(sp + 8)
            if isinstance(o, dict):
                o.pop(S.load_string(sp + 16), None)

        def valueIndex(sp):
            S.store_value(sp + 24, js_get(S.load_value(sp + 8), S.i64(sp + 16)))

        def valueSetIndex(sp):
            o, i, v = S.load_value(sp + 8), S.i64(sp + 16), S.load_value(sp + 24)
            if isinstance(o, list):
                while len(o) <= i:
                    o.append(UNDEF)
                o[i] = v
            elif isinstance(o, bytearray):
                o[i] = int(v) & 0xFF

        def call_common(fn, this, args, out):
            try:
                if not isinstance(fn, JSFunc):
                    raise TypeError(f"not a function: {fn!r}")
                r, ok = fn(this, *args), 1
            except Exception as e:  # noqa: BLE001
                S.log("[js throw]", e)
                r, ok = JSObject(message=str(e)), 0
            sp2 = S.sp()
            S.store_value(sp2 + out, r)
            S.wr(sp2 + out + 8, bytes([ok]))

        def valueCall(sp):
            v = S.load_value(sp + 8)
            m = js_get(v, S.load_string(sp + 16))
            call_common(m, v, S.load_values(sp + 32), 56)

        def valueInvoke(sp):
            call_common(S.load_value(sp + 8), UNDEF, S.load_values(sp + 16), 40)

        def valueNew(sp):
            call_common(S.load_value(sp + 8), UNDEF, S.load_values(sp + 16), 40)

        def valueLength(sp):
            S.set_i64(sp + 16, len(S.load_value(sp + 8)))

        def prepString(sp):
            b = js_str(S.load_value(sp + 8)).encode()
            S.store_value(sp + 16, U8Array(b))
            S.set_i64(sp + 24, len(b))

        def loadString(sp):
            s = S.load_value(sp + 8)
            a, n = S.load_slice(sp + 16)
            S.wr(a, bytes(s)[:n])

        def instanceOf(sp):
            S.wr(sp + 24, b"\x00")

        def toGo(sp):
            a, n = S.load_slice(sp + 8)
            src = S.load_value(sp + 32)
            if not isinstance(src, bytearray):
                S.wr(sp + 48, b"\x00")
                return
            c = bytes(src[:n])
            S.wr(a, c)
            S.set_i64(sp + 40, len(c))
            S.wr(sp + 48, b"\x01")

        def toJS(sp):
            dst = S.load_value(sp + 8)
            a, n = S.load_slice(sp + 16)
            if not isinstance(dst, bytearray):
                S.wr(sp + 48, b"\x00")
                return
            c = S.rd(a, min(n, len(dst)))
            dst[:len(c)] = c
            S.set_i64(sp + 40, len(c))
            S.wr(sp + 48, b"\x01")

        table = {
            "runtime.wasmExit": wasmExit, "runtime.wasmWrite": wasmWrite,
            "runtime.resetMemoryDataView": lambda sp: None, "runtime.nanotime1": nanotime1,
            "runtime.walltime": walltime, "runtime.scheduleTimeoutEvent": schedule,
            "runtime.clearTimeoutEvent": clear, "runtime.getRandomData": rand,
            "syscall/js.finalizeRef": finalize, "syscall/js.stringVal": stringVal,
            "syscall/js.valueGet": valueGet, "syscall/js.valueSet": valueSet,
            "syscall/js.valueDelete": valueDelete, "syscall/js.valueIndex": valueIndex,
            "syscall/js.valueSetIndex": valueSetIndex, "syscall/js.valueCall": valueCall,
            "syscall/js.valueInvoke": valueInvoke, "syscall/js.valueNew": valueNew,
            "syscall/js.valueLength": valueLength, "syscall/js.valuePrepareString": prepString,
            "syscall/js.valueLoadString": loadString, "syscall/js.valueInstanceOf": instanceOf,
            "syscall/js.copyBytesToGo": toGo, "syscall/js.copyBytesToJS": toJS,
            "debug": lambda v: S.log("[debug]", v),
        }
        imports = []
        for imp in self.module.imports:
            fn = table.get(imp.name)
            if fn is None:
                raise KeyError(imp.name)
            imports.append(mk(fn))
        return imports

    # ---------------- lifecycle
    def run(self):
        self.instance = wasmtime.Instance(self.store, self.module, self._imports())
        self.exports = self.instance.exports(self.store)
        self.memory = self.exports["mem"]
        off = 4096
        ptrs = []
        for s in ("js",):
            b = (s + "\0").encode()
            self.wr(off, b)
            ptrs.append(off)
            off += len(b)
            off += (-off) % 8
        ptrs += [0, 0]
        argv = off
        for p in ptrs:
            self.wr(off, struct.pack("<II", p, 0))
            off += 8
        self.exports["run"](self.store, 1, argv)
        self._drain_timeouts()

    def resume(self):
        self.exports["resume"](self.store)
        self._drain_timeouts()

    def _drain_timeouts(self, limit=1000):
        n = 0
        while self.timeouts and not self.exited and n < limit:
            tid = min(self.timeouts)
            self.timeouts.pop(tid)
            self.exports["resume"](self.store)
            n += 1

    def call_global(self, name, *args):
        fn = self.global_.get(name)
        if not isinstance(fn, JSFunc):
            raise KeyError(f"global {name} not registered: {fn!r}")
        return fn(UNDEF, *args)