agentmc/mcfold/anvil.py

"""从 MC 存档（anvil region）读真实区块数据。

为什么走存档而不走服务端：
  用户口径是「只调用渲染能力，不参与游戏」——存档是**零网络、零登录、
  零服务端**的真实世界数据源，且与官方客户端认同一套格式。

★ 本模块只做**读取**（Python 的职责）：
    region/*.mca → 解压区块 NBT → 取出 Sections[].block_states（palette + data）

  编码/协议组装**不在这里** —— 那属于 Java 侧（用官方 API 完成），
  因为历代改过协议格式（位宽、palette 分支、LongArray 编码），
  在 Python 里复制一份就等于把版本差异搬回了 Python。
  下游 render.py 只把读取结果压成 RLE，交给后端。
"""
from __future__ import annotations

import gzip
import io
import struct
import zlib
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator

SECTOR = 4096


class AnvilError(RuntimeError):
    pass


# ---------------------------------------------------------------- NBT 读取

class NbtReader:
    """极简 NBT 读取器（只读，支持存档里的全部 tag 类型）。"""

    def __init__(self, data: bytes):
        self.d = data
        self.i = 0

    def u1(self) -> int:
        v = self.d[self.i]; self.i += 1; return v

    def i1(self) -> int:
        v = struct.unpack_from(">b", self.d, self.i)[0]; self.i += 1; return v

    def i2(self) -> int:
        v = struct.unpack_from(">h", self.d, self.i)[0]; self.i += 2; return v

    def u2(self) -> int:
        v = struct.unpack_from(">H", self.d, self.i)[0]; self.i += 2; return v

    def i4(self) -> int:
        v = struct.unpack_from(">i", self.d, self.i)[0]; self.i += 4; return v

    def i8(self) -> int:
        v = struct.unpack_from(">q", self.d, self.i)[0]; self.i += 8; return v

    def f4(self) -> float:
        v = struct.unpack_from(">f", self.d, self.i)[0]; self.i += 4; return v

    def f8(self) -> float:
        v = struct.unpack_from(">d", self.d, self.i)[0]; self.i += 8; return v

    def s(self) -> str:
        n = self.u2()
        v = self.d[self.i:self.i + n].decode("utf-8", "replace"); self.i += n
        return v

    def payload(self, t: int) -> Any:
        if t == 1: return self.i1()
        if t == 2: return self.i2()
        if t == 3: return self.i4()
        if t == 4: return self.i8()
        if t == 5: return self.f4()
        if t == 6: return self.f8()
        if t == 7:
            n = self.i4(); v = self.d[self.i:self.i + n]; self.i += n; return v
        if t == 8: return self.s()
        if t == 9:
            it = self.u1(); n = self.i4()
            return [self.payload(it) for _ in range(n)]
        if t == 10:
            out = {}
            while True:
                tt = self.u1()
                if tt == 0: break
                nm = self.s()
                out[nm] = self.payload(tt)
            return out
        if t == 11:
            n = self.i4()
            v = list(struct.unpack_from(f">{n}i", self.d, self.i)); self.i += 4 * n
            return v
        if t == 12:
            n = self.i4()
            v = list(struct.unpack_from(f">{n}q", self.d, self.i)); self.i += 8 * n
            return v
        raise AnvilError(f"未知 NBT tag {t} @ {self.i}")


def nbt_read(data: bytes) -> dict:
    r = NbtReader(data)
    t = r.u1()
    if t != 10:
        raise AnvilError(f"根 tag 应为 Compound(10)，得到 {t}")
    r.s()                       # 根名（通常空串）
    return r.payload(10)


def read_level_dat(save_dir: Path) -> dict:
    """读 level.dat。

    ★ 实测两种布局都存在：
      - 官方启动器/服务端：8 字节头 + zlib
      - 部分整合包（本机 mc-test 实测）：文件直接就是 gzip
    所以先看魔数，再决定偏移。
    """
    raw = (save_dir / "level.dat").read_bytes()
    if raw[:2] == b"\x1f\x8b":                       # 直接 gzip
        body = gzip.decompress(raw)
    elif len(raw) > 10 and raw[8:10] == b"\x1f\x8b":  # 8 字节头 + gzip
        body = gzip.decompress(raw[8:])
    else:                                              # 8 字节头 + zlib
        body = zlib.decompress(raw[8:])
    return nbt_read(body)


# ---------------------------------------------------------------- region 读取

@dataclass
class RegionFile:
    """一个 .mca 文件：读出指定区块的原始 NBT（已解压）。"""

    path: Path
    _buf: bytes = b""

    def __post_init__(self):
        self._buf = self.path.read_bytes()
        if len(self._buf) < SECTOR * 2:
            raise AnvilError(f"region 文件过小: {self.path}")

    def chunk_nbt(self, cx: int, cz: int) -> bytes | None:
        """cx/cz 是**区块坐标**（世界坐标 >> 4）。"""
        lx, lz = cx & 31, cz & 31
        idx = (lx + lz * 32) * 4
        off = struct.unpack_from(">I", self._buf, idx)[0]
        if off == 0:
            return None
        sector = (off >> 8) & 0xFFFFFF
        count = off & 0xFF
        if sector == 0 or count == 0:
            return None
        start = sector * SECTOR
        length = struct.unpack_from(">I", self._buf, start)[0]
        comp = self._buf[start + 4]
        blob = self._buf[start + 5:start + 4 + length]
        if comp == 1:
            return gzip.decompress(blob)
        if comp == 2:
            return zlib.decompress(blob)
        if comp == 3:
            return blob
        if comp == 4:
            import lzma
            return lzma.decompress(blob)
        raise AnvilError(f"未知压缩类型 {comp}")


def region_path(save_dir: Path, dim: str | None, rx: int, rz: int) -> Path:
    dim_part = "" if not dim else f"{dim}/"
    return save_dir / f"{dim_part}region" / f"r.{rx}.{rz}.mca"


def list_region_files(save_dir: Path, dim: str | None = None) -> list[Path]:
    base = save_dir / (f"{dim}/region" if dim else "region")
    return sorted(base.glob("r.*.*.mca")) if base.is_dir() else []