"""protodef 解释器：读声明式 protocol.json，通用地编解码任意包。

为什么这不是"重写协议"（关键判断，2026-09-25 实测）：
    protocol.json 是**声明式**的 —— `varint`/`pstring`/`mapper`/`switch`
    全是数据描述而非代码。所以 111 个 play 包由同一份 JSON 驱动，
    我们只需实现 ~12 个原语 + 6 个复合类型，**不需要为任何单个包写代码**。

对照 mineflayer 侧：这对应 protodef（2045 行 JS）。Python 版可显著更短。
"""

from __future__ import annotations

import json
import struct
from pathlib import Path
from typing import Any


class ProtoError(Exception):
    pass


class Reader:
    """按协议原语从 buffer 顺序读取。"""

    __slots__ = ("buf", "pos")

    def __init__(self, buf: bytes, pos: int = 0):
        self.buf = buf
        self.pos = pos

    def remaining(self) -> int:
        return len(self.buf) - self.pos

    def read(self, n: int) -> bytes:
        if self.pos + n > len(self.buf):
            raise ProtoError(f"读越界: 需 {n} 字节, 剩 {self.remaining()}")
        b = self.buf[self.pos:self.pos + n]
        self.pos += n
        return b

    def varint(self) -> int:
        """LEB128 有符号变长整数（MC 的 varint）。"""
        result = 0
        shift = 0
        while True:
            b = self.read(1)[0]
            result |= (b & 0x7F) << shift
            if not (b & 0x80):
                break
            shift += 7
            if shift > 35:
                raise ProtoError("varint 过长")
        # 32 位有符号回绕
        if result & (1 << 31):
            result -= 1 << 32
        return result

    def varlong(self) -> int:
        result = 0
        shift = 0
        while True:
            b = self.read(1)[0]
            result |= (b & 0x7F) << shift
            if not (b & 0x80):
                break
            shift += 7
            if shift > 70:
                raise ProtoError("varlong 过长")
        if result & (1 << 63):
            result -= 1 << 64
        return result

    def pstring(self) -> str:
        n = self.varint()
        if n < 0:
            raise ProtoError(f"pstring 长度为负: {n}")
        return self.read(n).decode("utf-8", errors="replace")


class Writer:
    """按协议原语向 buffer 顺序写入。"""

    __slots__ = ("parts",)

    def __init__(self):
        self.parts: list[bytes] = []

    def getvalue(self) -> bytes:
        return b"".join(self.parts)

    def write(self, b: bytes) -> None:
        self.parts.append(b)

    def varint(self, v: int) -> None:
        if v < 0:
            v += 1 << 32
        while True:
            b = v & 0x7F
            v >>= 7
            if v:
                self.parts.append(bytes([b | 0x80]))
            else:
                self.parts.append(bytes([b]))
                return

    def pstring(self, s: str) -> None:
        b = s.encode("utf-8")
        self.varint(len(b))
        self.parts.append(b)

    # ---- 固定宽度原语（配置阶段要发 i8/u8/bool，实测标了 signed/unsigned） ----
    def _fixed(self, fmt: str, v: int) -> None:
        import struct as _s
        self.parts.append(_s.pack(fmt, v))

    def int8(self, v: int) -> None:
        """i8：有符号 8 位（如 viewDistance）。"""
        self._fixed(">b", v)

    def uint8(self, v: int) -> None:
        """u8：无符号 8 位（如 skinParts 位掩码 0x7F）。"""
        self._fixed(">B", v)

    def boolean(self, v: bool) -> None:
        """bool：协议里是 1 字节 0/1。"""
        self._fixed(">B", 1 if v else 0)

    def byte(self, v: int) -> None:
        """单字节（如 resource_pack 回应里的结果码 0..3）。"""
        self._fixed(">B", v)

    def int64(self, v: int) -> None:
        """i64：8 字节大端有符号（如 keep_alive 的 keepAliveId）。"""
        self._fixed(">q", v)


# ---- 固定宽度原语：[struct 格式, 字节数] ----
FIXED: dict[str, tuple[str, int]] = {
    "u8": (">B", 1), "i8": (">b", 1),
    "u16": (">H", 2), "i16": (">h", 2),
    "u32": (">I", 4), "i32": (">i", 4),
    "u64": (">Q", 8), "i64": (">q", 8),
    "f32": (">f", 4), "f64": (">d", 8),
    "bool": (">B", 1),
}


class ProtoDef:
    """按 protocol.json 的 types 表通用编解码。"""

    def __init__(self, types: dict[str, Any], phase_types: dict[str, Any] | None = None):
        # types: 全局 types；phase_types: 当前阶段作用域（如 play.toClient.types）
        # protodef 的具名类型会在阶段作用域里覆盖全局，必须两者都能查到。
        self.global_types = types
        self.phase_types = phase_types or {}
        self.types = {**types, **self.phase_types}

    # ---------- 读 ----------
    def read(self, t: Any, r: Reader, ctx: dict[str, Any] | None = None) -> Any:
        ctx = ctx if ctx is not None else {}

        # 1) 具名类型（含自定义类型展开）
        if isinstance(t, str):
            return self._read_named(t, r, ctx)

        if not isinstance(t, list) or not t:
            raise ProtoError(f"无法识别的类型声明: {t!r}")

        kind = t[0] if isinstance(t[0], str) else None

        # 2) 复合类型
        if kind == "container":
            out: dict[str, Any] = {}
            # ★ 子作用域：`../x` 路径要能回到外层（实测 6 处 `../flags/…`、
            #   `../action/…` 判据）。保留 __parent__ 链即可。
            sub = dict(ctx)
            sub["__parent__"] = ctx
            for f in t[1]:
                name = f.get("name") or f.get("anon")
                val = self.read(f["type"], r, sub)
                if f.get("anon"):
                    # 匿名字段：若是 dict 就并入
                    if isinstance(val, dict):
                        out.update(val)
                        sub.update(val)
                else:
                    out[name] = val
                    sub[name] = val
            # 回写外层（同名字段让外层也能查到，兼容平铺写法）
            for k, v in out.items():
                if k not in ("__parent__",):
                    ctx[k] = v
            return out

        if kind == "array":
            spec = t[1]
            if isinstance(spec, dict):
                # ★ protodef 的计数有两种写法（实测全协议：6 处 count / 103 处 countType）：
                #     countType: 'varint'            → 从流里读一个计数
                #     count: 'gridWidth'             → **引用已读字段**（ctx），不读流
                #   我先前只认 countType，遇到 `count` 就默认按 varint 读 ⇒
                #   多读一个字节 ⇒ **流错位**（表现为后面 switch 收到 948
                #   这种荒谬值，报“switch 无匹配分支”）。
                if "count" in spec and isinstance(spec["count"], str):
                    cname = spec["count"]
                    cnt = ctx.get(cname)
                    if cnt is None:
                        # 退化：当它是个要读的字段名
                        cnt = self.read(cname, r, ctx)
                    cnt = int(cnt or 0)
                else:
                    cnt = self._count(spec.get("countType", "varint"), r, ctx)
                if cnt < 0:
                    cnt = 0
                return [self.read(spec["type"], r, ctx) for _ in range(cnt)]
            # ["array", "u8"] 形式：读到结尾
            out = []
            while r.remaining() > 0:
                out.append(self.read(spec, r, ctx))
            return out

        if kind == "switch":
            spec = t[1]
            cmp_raw = spec.get("compareTo", "")
            # ★★ compareTo 是**路径表达式**，不一定是字段名（实测全协议）：
            #     "type"                        普通字段
            #     "flags/has_redirect_node"      子字段（跳进已读的 dict）
            #     "../flags/min_present"         父作用域的子字段
            #     "../action/add_player"         父作用域的位标志
            #   我先前只做 ctx.get(平铺名) ⇒ 这些全部查不到，
            #   退回 self.read('flags/xxx') ⇒ 报「未知类型: flags/xxx」
            #   （这正是 player_info / declare_commands 解不开的原因）。
            key = None
            val = None
            if isinstance(cmp_raw, str):
                key = cmp_raw[1:] if cmp_raw.startswith("$") else cmp_raw
                val = self._resolve_ctx(key, ctx)
            if val is None:
                # 退化：当它是个要读的类型名（罕见的隐式 switch）
                val = self.read(cmp_raw, r, ctx)

            fields = spec.get("fields", {})
            default = spec.get("default")
            # 字段键可能是 '0x01'/'1'/'true' 等写法，统一规整
            chosen = fields.get(str(val))
            if chosen is None and isinstance(val, bool):
                chosen = fields.get("1" if val else "0")
            if chosen is None:
                chosen = default
            if chosen is None:
                # 无匹配且无 default：按 void（不消费）
                if spec.get("default") == "void":
                    return None
                raise ProtoError(f"switch 无匹配分支: {key}={val!r}")
            return self.read(chosen, r, ctx)

        if kind == "option":
            present = self.read("bool", r, ctx)
            return self.read(t[1], r, ctx) if present else None

        if kind == "buffer":
            spec = t[1] if len(t) > 1 else {}
            cnt = self._count(spec.get("countType", "varint"), r, ctx) if spec else r.remaining()
            return r.read(cnt)

        if kind == "bitfield":
            total = sum(f["size"] for f in t[1])
            nbytes = (total + 7) // 8
            raw = int.from_bytes(r.read(nbytes), "big")
            out = {}
            shift = total
            for f in t[1]:
                shift -= f["size"]
                v = (raw >> shift) & ((1 << f["size"]) - 1)
                if f.get("signed") and v >= (1 << (f["size"] - 1)):
                    v -= 1 << f["size"]
                out[f["name"]] = v
            return out

        if kind == "mapper":
            spec = t[1]
            raw = self.read(spec["type"], r, ctx)
            for k, v in spec.get("mappings", {}).items():
                if int(k, 0) == raw:
                    return v
            return raw

        if kind == "void":
            return None

        # 列表形式引用内置原语，如 ["pstring", {"countType":"varint"}]、["varint"]
        if kind in FIXED or kind in ("varint", "optvarint", "varlong", "pstring",
                                     "string", "restBuffer", "UUID", "void",
                                     "anonymousNbt", "nbt", "optionalNbt"):
            return self._read_named(kind, r, ctx)

        # ---- bitflags：单个整数拆成多个布尔位（实测定义：PositionUpdateRelatives）
        #
        # 参照实现 protodef/src/datatypes/utils.js readBitflags：
        #   返回 {_value: 原始值, <flag>: bool, ...}
        # ★ 为什么必须实现：`position` 包里的 `flags` 就是它。
        #   position 是「玩家在哪」的唯一来源 —— 实现不了它就等于
        #   看不到自己位置，后续 M3 状态层根本无从上手。
        if kind == "bitflags":
            spec = t[1]
            base = spec.get("type", "u8")
            raw = self.read(base, r, ctx)
            flags = spec.get("flags")
            out: dict[str, Any] = {"_value": raw}
            if isinstance(flags, list):
                for i, nm in enumerate(flags):
                    out[nm] = bool(raw & (1 << i))
            elif isinstance(flags, dict):
                shift = spec.get("shift")
                for nm, bit in flags.items():
                    b = shift if shift else bit
                    out[nm] = bool(raw & (1 << int(b)))
            return out

        # ---- bitfield：把一个整数的各位切成具名子字段（实测定义：command_node.flags）
        # ---- entityMetadataLoop：读实体元数据条目直到终止值（实测：entityMetadata）
        #
        # ★ 终止判据是**键名等于 endVal**（协议里 metadata 以 key=0xFF 收尾），
        #   不是读到流尽头 —— 写成读尽头会在同包多个 metadata 时错位。
        if kind == "entityMetadataLoop":
            spec = t[1]
            end_val = spec.get("endVal", 255)
            entry_t = spec.get("type")
            out_list: list[Any] = []
            # ★ 终止符是**光秃秃的一个字节**（0xFF），后面**没有** type/value。
            #   我先前「读完整个条目再看 key 是否 == endVal」⇒ 终止时多读了
            #   type+value 2 字节 ⇒ 报「读越界」。
            #   参照 minecraft-protocol/src/datatypes/minecraft.js:120-128：
            #     item = buffer.readUInt8(cursor); if (item === endVal) 停
            #   线上实测（entityId=144 的真实包 58 90 01 09 03 41a00000 11 00 7f ff）：
            #   按窥视语义切，13/13 字节完全对齐。
            while True:
                if r.remaining() <= 0:
                    break
                peek = r.buf[r.pos]
                if peek == end_val:
                    r.read(1)          # 只消费终止键本身
                    break
                start = r.pos
                entry = self.read(entry_t, r, ctx)
                out_list.append(entry)
                if r.pos == start:
                    break   # 防死循环（零字节推进）
            return out_list

        raise ProtoError(f"未实现的复合类型: {kind}")

    @staticmethod
    def _resolve_ctx(path: str, ctx: dict[str, Any]) -> Any:
        """按路径表达式取值：`a/b`、`../a/b`、`a.b` 都要支持。

        ★ 实测语法（全协议 12 处带路径的 compareTo）：
            flags/has_redirect_node         子字段
            ../flags/has_custom_suggestions 父作用域子字段
        ★ 为什么需“向上逐层试”：`..` 的层数依赖引用点嵌套深度，
          而不同分支的嵌套深度不一样 ⇒ 固定按字面层级走会在某些分支落空，
          落空后代码会退化成“读一个叫 ../flags/… 的类型”
          （报「未知类型: ../flags/has_custom_suggestions」）。
        """
        parts = [p for p in path.replace(".", "/").split("/") if p]
        ups = 0
        while parts and parts[0] == "..":
            parts.pop(0)
            ups += 1
        if not parts:
            return None
        # 先按字面层级走，失败则向上多跳，直到祖先链尽头
        for extra in range(0, 8):
            for total in ({ups + extra} if extra == 0 else {ups - extra, ups + extra}):
                if total < 0:
                    continue
                cur: Any = ctx
                for _ in range(total):
                    if not isinstance(cur, dict):
                        cur = None
                        break
                    cur = cur.get("__parent__")
                if cur is None:
                    continue
                ok = True
                for part in parts:
                    if isinstance(cur, dict) and part in cur:
                        cur = cur[part]
                    else:
                        ok = False
                        break
                if ok:
                    return cur
        return None

    def _read_named(self, name: str, r: Reader, ctx: dict[str, Any]) -> Any:
        # 阶段作用域优先（packet_set_protocol 这类只存在于阶段内）
        if name in self.phase_types:
            return self.read(self.phase_types[name], r, ctx)
        # protocol.json 里 "native" 表示「此类型由实现直接提供」，
        # 不是要展开的定义 —— 必须按内置原语处理（实测踩过这个坑）。
        if name in self.global_types and self.global_types[name] != "native":
            return self.read(self.global_types[name], r, ctx)

        if name in FIXED:
            fmt, n = FIXED[name]
            v = struct.unpack(fmt, r.read(n))[0]
            return bool(v) if name == "bool" else v

        if name == "varint":
            return r.varint()
        if name == "optvarint":
            return r.varint()
        if name == "varlong":
            return r.varlong()
        if name in ("string", "pstring"):
            return r.pstring()
        if name == "restBuffer":
            return r.read(r.remaining())
        if name == "void":
            return None
        if name == "UUID":
            b = r.read(16)
            h = b.hex()
            return f"{h[:8]}-{h[8:12]}-{h[12:16]}-{h[16:20]}-{h[20:]}"
        if name == "nbt":
            return self._read_nbt(r)
        if name == "anonymousNbt":
            return self._read_anonymous_nbt(r)
        if name == "optionalNbt":
            return self._read_nbt(r) if r.read(1)[0] != 0 else None

        raise ProtoError(f"未知类型: {name}")

    # ---- NBT ----
    # ★ 为什么必须真解 NBT（实测 2026-09-26）：
    #   configuration 阶段的 registry_data 条目是 `anonymousNbt` ——
    #   方块/物品注册表的**权威内容就在 NBT 里**（mod 方块的 state 表也在这）。
    #   旧版只回 {"_nbt_raw": ...} 占位，等于拿不到任何注册表内容，
    #   上层只能看到 `unknown(id)` —— 这正是「与服务端 mod 物品交互」的卡点。
    _NBT_END, _NBT_BYTE, _NBT_SHORT, _NBT_INT, _NBT_LONG = 0, 1, 2, 3, 4
    _NBT_FLOAT, _NBT_DOUBLE, _NBT_BYTE_ARRAY, _NBT_STRING = 5, 6, 7, 8
    _NBT_LIST, _NBT_COMPOUND, _NBT_INT_ARRAY, _NBT_LONG_ARRAY = 9, 10, 11, 12

    @classmethod
    def _read_nbt_payload(cls, r: Reader, tag: int, depth: int = 0) -> Any:
        """按 tag 号读一个 NBT 载荷（不含 tag 字节本身）。"""
        import struct as _s
        if depth > 64:
            raise ProtoError("NBT 嵌套过深")
        if tag == cls._NBT_BYTE:
            return r.read(1)[0]
        if tag == cls._NBT_SHORT:
            return int.from_bytes(r.read(2), "big", signed=True)
        if tag == cls._NBT_INT:
            return int.from_bytes(r.read(4), "big", signed=True)
        if tag == cls._NBT_LONG:
            return int.from_bytes(r.read(8), "big", signed=True)
        if tag == cls._NBT_FLOAT:
            return _s.unpack(">f", r.read(4))[0]
        if tag == cls._NBT_DOUBLE:
            return _s.unpack(">d", r.read(8))[0]
        if tag == cls._NBT_BYTE_ARRAY:
            n = int.from_bytes(r.read(4), "big", signed=True)
            return list(r.read(max(n, 0)))
        if tag == cls._NBT_STRING:
            n = int.from_bytes(r.read(2), "big")
            return r.read(n).decode("utf-8", "replace")
        if tag == cls._NBT_LIST:
            et = r.read(1)[0]
            n = int.from_bytes(r.read(4), "big", signed=True)
            if et == cls._NBT_END or n <= 0:
                return []
            return [cls._read_nbt_payload(r, et, depth + 1) for _ in range(n)]
        if tag == cls._NBT_COMPOUND:
            out: dict[str, Any] = {}
            while True:
                t = r.read(1)[0]
                if t == cls._NBT_END:
                    return out
                nm_len = int.from_bytes(r.read(2), "big")
                nm = r.read(nm_len).decode("utf-8", "replace")
                out[nm] = cls._read_nbt_payload(r, t, depth + 1)
        if tag in (cls._NBT_INT_ARRAY, cls._NBT_LONG_ARRAY):
            w = 4 if tag == cls._NBT_INT_ARRAY else 8
            n = int.from_bytes(r.read(4), "big", signed=True)
            return [int.from_bytes(r.read(w), "big", signed=True) for _ in range(max(n, 0))]
        raise ProtoError(f"未知 NBT tag: {tag}")

    @classmethod
    def _read_nbt(cls, r: Reader) -> Any:
        """读一个具名 NBT（tag + 名字 + 载荷）—— 未压缩的普通 NBT。"""
        tag = r.read(1)[0]
        if tag == cls._NBT_END:
            return None
        n = int.from_bytes(r.read(2), "big")
        name = r.read(n).decode("utf-8", "replace")
        return {"name": name, "value": cls._read_nbt_payload(r, tag)}

    @classmethod
    def _read_anonymous_nbt(cls, r: Reader) -> Any:
        """anonymousNbt：无根名字的 NBT（网络格式，registry_data 用的就是这个）。

        ★ 与普通 NBT 的差别只在**没有根名字** —— 少读一个
          (short 长度 + 字节) 就直接进载荷。读多了会把后面
          所有字节都错位（并且不报错，只是字段全乱）。
        """
        tag = r.read(1)[0]
        if tag == cls._NBT_END:
            return None
        return cls._read_nbt_payload(r, tag)

    def _count(self, count_type: Any, r: Reader, ctx: dict[str, Any]) -> int:
        return self.read(count_type, r, ctx)


def for_phase(protocol: dict[str, Any], phase: str, side: str) -> ProtoDef:
    """按「阶段+方向」构造解释器（如 play/toClient）。

    这是正确用法：protodef 的具名类型分全局与阶段两级作用域，
    只传全局表会导致 packet_set_protocol 这类阶段内类型查不到。
    """
    ph = protocol.get(phase, {}).get(side, {})
    return ProtoDef(protocol.get("types", {}), ph.get("types", {}))


def load_protocol(path: str | Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)