import ctypes
from pathlib import Path

from wasmtime import Engine, Linker, Module, Store, WasiConfig

_WASM_PATH = Path(__file__).parent / "_vt100.wasm"

_engine = Engine()
_module = Module.from_file(_engine, str(_WASM_PATH))


def _unpack_result(packed: int) -> tuple[int, int]:
    ptr = (packed >> 32) & 0xFFFFFFFF
    length = packed & 0xFFFFFFFF
    return ptr, length


class _Runtime:
    def __init__(self) -> None:
        self._store = Store(_engine)
        config = WasiConfig()
        self._store.set_wasi(config)
        linker = Linker(_engine)
        linker.define_wasi()
        instance = linker.instantiate(self._store, _module)
        exports = instance.exports(self._store)
        self._memory = exports["memory"]
        self._alloc_fn = exports["wasm_alloc"]
        self._dealloc_fn = exports["wasm_dealloc"]
        self._apply_delta_fn = exports["apply_delta"]
        self._screen_contents_fn = exports["screen_contents"]
        self._screen_contents_formatted_fn = exports["screen_contents_formatted"]

    def _write_bytes(self, data: bytes) -> tuple[int, int]:
        size = len(data)
        if size == 0:
            return 0, 0
        ptr = self._alloc_fn(self._store, size)
        if ptr == 0:
            raise MemoryError(f"WASM allocation failed for {size} bytes")
        self._memory.write(self._store, data, ptr)
        return ptr, size

    def _read_result(self, packed: int) -> bytes:
        ptr, length = _unpack_result(packed)
        if ptr == 0 or length == 0:
            return b""
        base = ctypes.addressof(self._memory.data_ptr(self._store).contents)
        result = (ctypes.c_char * length).from_address(base + ptr).raw
        self._dealloc_fn(self._store, ptr, length)
        return result

    def apply_delta(self, state: bytes | None, delta: bytes, rows: int, cols: int) -> bytes:
        if state is not None:
            s_ptr, s_len = self._write_bytes(state)
        else:
            s_ptr, s_len = 0, 0
        d_ptr, d_len = self._write_bytes(delta)
        packed = self._apply_delta_fn(self._store, s_ptr, s_len, d_ptr, d_len, rows, cols)
        if s_ptr:
            self._dealloc_fn(self._store, s_ptr, s_len)
        if d_ptr:
            self._dealloc_fn(self._store, d_ptr, d_len)
        return self._read_result(packed)

    def screen_contents(self, state: bytes) -> str:
        s_ptr, s_len = self._write_bytes(state)
        packed = self._screen_contents_fn(self._store, s_ptr, s_len)
        if s_ptr:
            self._dealloc_fn(self._store, s_ptr, s_len)
        return self._read_result(packed).decode("utf-8")

    def screen_contents_formatted(self, state: bytes) -> bytes:
        s_ptr, s_len = self._write_bytes(state)
        packed = self._screen_contents_formatted_fn(self._store, s_ptr, s_len)
        if s_ptr:
            self._dealloc_fn(self._store, s_ptr, s_len)
        return self._read_result(packed)


_runtime = _Runtime()


def apply_delta(state: bytes | None, delta: bytes, rows: int, cols: int) -> bytes:
    return _runtime.apply_delta(state, delta, rows, cols)


def screen_contents(state: bytes) -> str:
    return _runtime.screen_contents(state)


def screen_contents_formatted(state: bytes) -> bytes:
    return _runtime.screen_contents_formatted(state)


__all__ = ["apply_delta", "screen_contents", "screen_contents_formatted"]
