import ctypes

from . import _ffi as ffi
from wasmtime import TableType, Store, WasmtimeError, Val
from typing import Optional, Any
from ._store import Storelike


class Table:
    _table: ffi.wasmtime_table_t

    def __init__(self, store: Store, ty: TableType, init: Any):
        """
        Creates a new table within `store` with the specified `ty`.
        """

        init_val = Val._convert_to_raw(store, ty.element, init)

        table = ffi.wasmtime_table_t()
        error = ffi.wasmtime_table_new(store._context(), ty.ptr(), ctypes.byref(init_val), ctypes.byref(table))
        ffi.wasmtime_val_unroot(ctypes.byref(init_val))
        if error:
            raise WasmtimeError._from_ptr(error)
        self._table = table

    @classmethod
    def _from_raw(cls, table: ffi.wasmtime_table_t) -> "Table":
        ty: "Table" = cls.__new__(cls)
        ty._table = table
        return ty

    def type(self, store: Storelike) -> TableType:
        """
        Gets the type of this table as a `TableType`
        """

        ptr = ffi.wasmtime_table_type(store._context(), ctypes.byref(self._table))
        return TableType._from_ptr(ptr, None)

    def size(self, store: Storelike) -> int:
        """
        Gets the size, in elements, of this table
        """
        return ffi.wasmtime_table_size(store._context(), ctypes.byref(self._table))

    def grow(self, store: Storelike, amt: int, init: Any) -> int:
        """
        Grows this table by the specified number of slots, using the specified
        initializer for all new table slots.

        Raises a `WasmtimeError` if the table could not be grown.
        Returns the previous size of the table otherwise.
        """
        init_val = Val._convert_to_raw(store, self.type(store).element, init)
        prev = ctypes.c_uint64(0)
        error = ffi.wasmtime_table_grow(store._context(), ctypes.byref(self._table), ctypes.c_uint64(amt), ctypes.byref(init_val), ctypes.byref(prev))
        ffi.wasmtime_val_unroot(ctypes.byref(init_val))
        if error:
            raise WasmtimeError._from_ptr(error)
        return prev.value

    def get(self, store: Store, idx: int) -> Optional[Any]:
        """
        Gets an individual element within this table.

        Returns `None` for null references in the table (i.e. a null `funcref`
        or a null `externref).

        Returns a `Func` for non-null `funcref` table elements.

        Returns the wrapped extern data for non-null `externref` table elements.

        Returns `None` if `idx` is out of bounds.
        """
        raw = ffi.wasmtime_val_t()
        ok = ffi.wasmtime_table_get(store._context(), ctypes.byref(self._table), idx, ctypes.byref(raw))
        if not ok:
            return None
        val = Val._from_raw(store, raw)
        if val.value:
            return val.value
        else:
            return val

    def set(self, store: Store, idx: int, val: Any) -> None:
        """
        Sets an individual element within this table.

        `idx` must be an integer index.

        The `val` specified must be convertible into this table's element
        type. I.e. for a `funcref` table, `val` must either be a `Func` or
        `None`, and for an `externref` table, `val` may be any arbitrary
        external data.

        Raises a `WasmtimeError` if `idx` is out of bounds.
        """
        value = Val._convert_to_raw(store, self.type(store).element, val)
        error = ffi.wasmtime_table_set(store._context(), ctypes.byref(self._table), idx, ctypes.byref(value))
        ffi.wasmtime_val_unroot(ctypes.byref(value))
        if error:
            raise WasmtimeError._from_ptr(error)

    def _as_extern(self) -> ffi.wasmtime_extern_t:
        union = ffi.wasmtime_extern_union(table=self._table)
        return ffi.wasmtime_extern_t(ffi.WASMTIME_EXTERN_TABLE, union)
