import t, c
import memhub
import string
import json
from stdint import *

# Variant 类型标签 (与 _variant.py 保持一致)
VARIANT_NONE:   t.CDefine = 0   # 无类型/空值
VARIANT_INT:    t.CDefine = 1   # int (i64)
VARIANT_FLOAT:  t.CDefine = 2   # float (f32)
VARIANT_DOUBLE: t.CDefine = 3   # double (f64)
VARIANT_STR:    t.CDefine = 4   # str (i8*)
VARIANT_DICT:   t.CDefine = 5   # dict*
VARIANT_LIST:   t.CDefine = 6   # list*
VARIANT_BOOL:   t.CDefine = 7   # bool (i1/i32)
VARIANT_PTR:    t.CDefine = 8   # 自定义指针类型

# Variant 结构体大小: vtype(i32, 4 bytes) + padding(4 bytes) + ptr(i8*, 8 bytes) = 16 bytes
_VARIANT_SIZE: t.CDefine = 16


class dict:
    """堆上动态字典容器，基于 mbuddy 分配器。
    键类型固定为 str (i8*)，值使用 Variant 存储 (类型标签 + 无类型指针)。
    支持嵌套 dict/list 和 JSON 序列化。

    用法:
        mb: memhub.MemBuddy = memhub.MemBuddy(arena, arena_size)
        d: dict = dict(mb)
        d.set_int("x", 42)
        p: t.CPtr = d["x"]
        v: int = c.Deref(p)

        # 嵌套 dict
        inner: dict = dict(mb)
        inner.set_int("a", 1)
        d.set_dict("sub", inner)
        sub: dict = c.Deref(d["sub"])
        av: int = c.Deref(sub["a"])
    """
    __keys__:       t.CPtr
    __values__:     t.CPtr
    __count__:      t.CSizeT = 0
    __capacity__:   t.CSizeT = 8
    __mbuddy__:     memhub.MemBuddy | t.CPtr
    __iter_index__: t.CSizeT = 0

    def __new__(self, mb: memhub.MemBuddy | t.CPtr) -> t.CPtr:
        # 堆分配 dict 对象
        return mb.alloc(dict.__sizeof__())

    def __init__(self, mb: memhub.MemBuddy | t.CPtr):
        self.__mbuddy__ = mb
        self.__keys__   = mb.alloc(self.__capacity__ * 8)
        self.__values__ = mb.alloc(self.__capacity__ * _VARIANT_SIZE)

    def __len__(self) -> t.CSizeT:
        return self.__count__

    def _find(self, key: str) -> t.CSizeT:
        """查找 key 的索引，找不到返回 __count__"""
        for i in range(self.__count__):
            key_ptr: str | t.CPtr = t.CPtr(t.CUInt64T(self.__keys__) + i * 8)
            if key_ptr[0] == key: return i
        return self.__count__

    def __getitem__(self, key: str) -> t.CPtr:
        """返回值的无类型指针。用户用 c.Deref 解引用获取实际值。
        例: p: t.CPtr = d["x"]; v: int = c.Deref(p)"""
        idx: t.CSizeT = self._find(key)
        if idx >= self.__count__: return None
        # Variant.ptr 在 Variant 结构体偏移 8 处
        val_ptr: str | t.CPtr = t.CPtr(t.CUInt64T(self.__values__) + idx * _VARIANT_SIZE + 8)
        return val_ptr[0]

    def __setitem__(self, key: str, val: t.CNeedPtr):
        """存储值为 VARIANT_PTR。val 自动取地址 (t.CNeedPtr)。
        注意: val 指向的值生命周期由调用者管理。
        如需持久存储，请使用 set_int / set_str 等类型化方法。"""
        self._set_variant_raw(key, VARIANT_PTR, val)

    def _set_variant_raw(self, key: str, vtype: t.CInt, ptr: t.CPtr):
        """内部方法: 存储 Variant (vtype + ptr)。不自动取地址。"""
        idx: t.CSizeT = self._find(key)
        if idx < self.__count__:
            # 更新已有键
            vtype_ptr: t.CInt | t.CPtr = t.CPtr(t.CUInt64T(self.__values__) + idx * _VARIANT_SIZE)
            vtype_ptr[0] = vtype
            val_ptr: str | t.CPtr = t.CPtr(t.CUInt64T(self.__values__) + idx * _VARIANT_SIZE + 8)
            val_ptr[0] = ptr
            return
        if self.__count__ >= self.__capacity__:
            new_cap: t.CSizeT = self.__capacity__ * 2
            new_keys: t.CPtr = self.__mbuddy__.alloc(new_cap * 8)
            new_values: t.CPtr = self.__mbuddy__.alloc(new_cap * _VARIANT_SIZE)
            if new_keys is None or new_values is None: return
            string.memcpy(new_keys,   self.__keys__,   self.__count__ * 8)
            string.memcpy(new_values, self.__values__, self.__count__ * _VARIANT_SIZE)
            self.__keys__     = new_keys
            self.__values__   = new_values
            self.__capacity__ = new_cap
        key_ptr: str | t.CPtr = t.CPtr(t.CUInt64T(self.__keys__) + self.__count__ * 8)
        key_ptr[0] = key
        vtype_ptr: t.CInt | t.CPtr = t.CPtr(t.CUInt64T(self.__values__) + self.__count__ * _VARIANT_SIZE)
        vtype_ptr[0] = vtype
        val_ptr: str | t.CPtr = t.CPtr(t.CUInt64T(self.__values__) + self.__count__ * _VARIANT_SIZE + 8)
        val_ptr[0] = ptr
        self.__count__ += 1

    # ============================================================
    # 类型化 setter — 在 __mbuddy__ 中分配持久存储
    # ============================================================

    def set_int(self, key: str, val: int):
        """存储 int 值 (VARIANT_INT)。值复制到 __mbuddy__ 持久存储。"""
        storage: t.CPtr = self.__mbuddy__.alloc(8)
        int_ptr: int | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
        int_ptr[0] = val
        self._set_variant_raw(key, VARIANT_INT, storage)

    def set_str(self, key: str, val: str):
        """存储 str 值 (VARIANT_STR)。直接存储字符串指针。"""
        self._set_variant_raw(key, VARIANT_STR, val)

    def set_float(self, key: str, val: float):
        """存储 float 值 (VARIANT_FLOAT)。值复制到 __mbuddy__ 持久存储。"""
        storage:   t.CPtr = self.__mbuddy__.alloc(8)
        float_ptr: float | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
        float_ptr[0] = val
        self._set_variant_raw(key, VARIANT_FLOAT, storage)

    def set_double(self, key: str, val: t.CDouble):
        """存储 double 值 (VARIANT_DOUBLE)。值复制到 __mbuddy__ 持久存储。"""
        storage: t.CPtr = self.__mbuddy__.alloc(8)
        dbl_ptr: t.CDouble | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
        dbl_ptr[0] = val
        self._set_variant_raw(key, VARIANT_DOUBLE, storage)

    def set_dict(self, key: str, val: dict | t.CPtr):
        """存储嵌套 dict 指针 (VARIANT_DICT)。"""
        self._set_variant_raw(key, VARIANT_DICT, val)

    def set_list(self, key: str, val: t.CPtr):
        """存储嵌套 list 指针 (VARIANT_LIST)。"""
        self._set_variant_raw(key, VARIANT_LIST, val)

    def set_bool(self, key: str, val: t.CInt):
        """存储 bool 值 (VARIANT_BOOL)。值复制到 __mbuddy__ 持久存储。"""
        storage: t.CPtr = self.__mbuddy__.alloc(8)
        bool_ptr: t.CInt | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
        bool_ptr[0] = val
        self._set_variant_raw(key, VARIANT_BOOL, storage)

    # ============================================================
    # 查询方法
    # ============================================================

    def get_vtype(self, key: str) -> t.CInt:
        """返回 key 对应值的 Variant 类型标签。"""
        idx: t.CSizeT = self._find(key)
        if idx >= self.__count__:
            return VARIANT_NONE
        vtype_ptr: t.CInt | t.CPtr = t.CVoid(t.CUInt64T(self.__values__) + idx * _VARIANT_SIZE, t.CPtr)
        return vtype_ptr[0]

    def get(self, key: str, default: t.CPtr) -> t.CPtr:
        """返回值指针，不存在则返回 default。"""
        idx: t.CSizeT = self._find(key)
        if idx >= self.__count__:
            return default
        val_ptr: str | t.CPtr = t.CVoid(t.CUInt64T(self.__values__) + idx * _VARIANT_SIZE + 8, t.CPtr)
        return val_ptr[0]

    def contains(self, key: str) -> t.CInt:
        """检查 key 是否存在。"""
        idx: t.CSizeT = self._find(key)
        if idx < self.__count__:
            return 1
        return 0

    def __iter__(self) -> dict | t.CPtr:
        self.__iter_index__ = 0
        return self

    def __next__(self) -> str:
        if self.__iter_index__ >= self.__count__:
            raise StopIteration
        idx: t.CSizeT = self.__iter_index__
        self.__iter_index__ = idx + 1
        key_ptr: str | t.CPtr = t.CVoid(t.CUInt64T(self.__keys__) + idx * 8, t.CPtr)
        return key_ptr[0]

    # ============================================================
    # JSON 处理：利用 includes/json 库实现序列化/反序列化
    # ============================================================

    def to_json(self, pool: memhub.MemManager | t.CPtr) -> json.JsonValue | t.CPtr:
        """将 dict 转换为 JsonValue 对象 (JSON_OBJECT)。
        根据 Variant 类型标签自动选择合适的 JSON 值类型。"""
        obj: json.JsonValue | t.CPtr = json.object(pool)
        i: t.CSizeT = 0
        while i < self.__count__:
            key_ptr: str | t.CPtr = t.CVoid(t.CUInt64T(self.__keys__) + i * 8, t.CPtr)
            vtype_ptr: t.CInt | t.CPtr = t.CVoid(t.CUInt64T(self.__values__) + i * _VARIANT_SIZE, t.CPtr)
            val_ptr: str | t.CPtr = t.CVoid(t.CUInt64T(self.__values__) + i * _VARIANT_SIZE + 8, t.CPtr)
            vt: t.CInt = vtype_ptr[0]
            vp: t.CPtr = val_ptr[0]
            jv: json.JsonValue | t.CPtr = json.null(pool)
            if jv != None:
                if vt == VARIANT_INT:
                    jv.vtype = 2  # JSON_INT
                    ip: int | t.CPtr = t.CVoid(t.CUInt64T(vp), t.CPtr)
                    jv.int_val = ip[0]
                elif vt == VARIANT_STR:
                    jv.vtype = 4  # JSON_STRING
                    jv.str_val = vp
                elif vt == VARIANT_BOOL:
                    jv.vtype = 1  # JSON_BOOL
                    bp: t.CInt | t.CPtr = t.CVoid(t.CUInt64T(vp), t.CPtr)
                    jv.bool_val = bp[0]
                elif vt == VARIANT_FLOAT:
                    jv.vtype = 3  # JSON_FLOAT
                    fp: float | t.CPtr = t.CVoid(t.CUInt64T(vp), t.CPtr)
                    jv.float_val = fp[0]
                elif vt == VARIANT_PTR:
                    # 默认: 将 PTR 作为整数处理
                    jv.vtype = 2  # JSON_INT
                    ip: int | t.CPtr = t.CVoid(t.CUInt64T(vp), t.CPtr)
                    jv.int_val = ip[0]
            json.object_set(pool, obj, key_ptr[0], jv)
            i += 1
        return obj

    def from_json(self, pool: memhub.MemManager | t.CPtr, root: json.JsonValue | t.CPtr):
        """从 JsonValue 对象填充 dict。
        根据 JSON 值类型自动提取并存储为对应 Variant 类型。"""
        if root == None:
            return
        if root.vtype != 6:  # JSON_OBJECT
            return
        cur: json.JsonValue | t.CPtr = root.child
        while cur != None:
            if cur.key != None:
                if cur.vtype == 2:  # JSON_INT
                    storage: t.CPtr = self.__mbuddy__.alloc(8)
                    ip: int | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
                    ip[0] = cur.int_val
                    self._set_variant_raw(cur.key, VARIANT_INT, storage)
                elif cur.vtype == 4:  # JSON_STRING
                    self._set_variant_raw(cur.key, VARIANT_STR, cur.str_val)
                elif cur.vtype == 1:  # JSON_BOOL
                    storage: t.CPtr = self.__mbuddy__.alloc(8)
                    bp: t.CInt | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
                    bp[0] = cur.bool_val
                    self._set_variant_raw(cur.key, VARIANT_BOOL, storage)
                elif cur.vtype == 3:  # JSON_FLOAT
                    storage: t.CPtr = self.__mbuddy__.alloc(8)
                    fp: float | t.CPtr = t.CVoid(t.CUInt64T(storage), t.CPtr)
                    fp[0] = cur.float_val
                    self._set_variant_raw(cur.key, VARIANT_FLOAT, storage)
            cur = cur.next

    def dumps(self, pool: memhub.MemManager | t.CPtr) -> str:
        """将 dict 序列化为 JSON 字符串。"""
        obj: json.JsonValue | t.CPtr = self.to_json(pool)
        return json.write(pool, obj, False)

    def loads(self, pool: memhub.MemManager | t.CPtr, json_str: str):
        """从 JSON 字符串解析并填充 dict。"""
        root: json.JsonValue | t.CPtr = json.parse(pool, json_str)
        self.from_json(pool, root)