"""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)


# ---- 固定宽度原语：[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] = {}
            for f in t[1]:
                name = f.get("name") or f.get("anon")
                val = self.read(f["type"], r, ctx)
                if f.get("anon"):
                    # 匿名字段：若是 dict 就并入
                    if isinstance(val, dict):
                        out.update(val)
                else:
                    out[name] = val
                    ctx[name] = val
            return out

        if kind == "array":
            spec = t[1]
            if isinstance(spec, dict):
                cnt = self._count(spec.get("countType", "varint"), r, ctx)
                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", "")
            # protodef 两种写法："$field" 或 "field"（都指引用已读字段）
            key = None
            if isinstance(cmp_raw, str):
                key = cmp_raw[1:] if cmp_raw.startswith("$") else cmp_raw
            val = ctx.get(key) if (key and key in ctx) else self.read(cmp_raw, r, ctx)

            fields = spec.get("fields", {})
            default = spec.get("default")
            # 字段键可能是 '0x01'/'1' 等写法，统一规整
            chosen = fields.get(str(val))
            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"):
            return self._read_named(kind, r, ctx)

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

    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 == "optionalNbt":
            return self._read_nbt(r) if r.read(1)[0] != 0 else None

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

    @staticmethod
    def _read_nbt(r: Reader) -> Any:
        """最小 NBT 读取（够用于 registry/codec 探测）。"""
        tag = r.read(1)[0]
        if tag == 0:
            return None
        return {"_nbt_tag": tag, "_nbt_raw": r.read(r.remaining())}

    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)