"""把世界数据渲染成画面 —— mcfold 的对外渲染入口。

分工（与设计文档附录 D 一致）：
  本模块（Python）  : 世界数据 —— 区块、方块、光照、相机
  客户端 jar（Java）: 方块名与 state id 的权威映射 + 画面

为什么 state id 必须由渲染器侧解析：
  存档 palette 里是**方块名**（`minecraft:oak_log[axis=y]`），
  而网络载荷要 **state id**。名字→id 的映射随版本/加载器/mod 变化，
  只有客户端自带的注册表是权威的（1.20.1 不由网络下发，附录 D）。

用法：
    from mcfold.render import WorldRenderer
    with WorldRenderer.open(instance) as wr:
        wr.load_save("/home/mc-test/mc-test", center=(-1, -1), radius=3)
        wr.camera(x=8, y=80, z=8, yaw=135, pitch=35)
        img = wr.screenshot("/tmp/out.png")
"""
from __future__ import annotations

import base64
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable

from . import anvil
from .render_bridge import RenderBridge, RenderBridgeError, RenderResult

# 存档区块在网络载荷里的 section 数（1.20.1 worldHeight 384 / 16 = 24）
SECTION_COUNT = 24
SECTION_VOLUME = 16 * 16 * 16
BIOME_VOLUME = 4 * 4 * 4
# ★ 最小 section 的 Y：1.20 起 worldHeight 384、minBuildHeight -64 ⇒ minSection = -4。
#   网络载荷里的 section 索引是「从下往上数」（0..23），而存档 NBT 用的是绝对 Y（-4..19）。
MIN_SECTION_Y = -4

# 位打包：把 palette 索引列表压进 long（值不跨 long 边界）
def _pack(values: list[int], bits: int) -> list[int]:
    if bits <= 0:
        return []
    vpl = 64 // bits
    mask = (1 << bits) - 1
    out, cur, slot = [], 0, 0
    for v in values:
        cur |= (v & mask) << (slot * bits)
        slot += 1
        if slot == vpl:
            out.append(cur); cur, slot = 0, 0
    if slot:
        out.append(cur)
    return out


def _varint(v: int) -> bytes:
    out = bytearray()
    while True:
        b = v & 0x7F; v >>= 7
        if v: out.append(b | 0x80)
        else: out.append(b); break
    return bytes(out)


def _long_array(longs: list[int]) -> bytes:
    out = _varint(len(longs))
    for l in longs:
        out += (l & 0xFFFFFFFFFFFFFFFF).to_bytes(8, "big")
    return out


def _encode_paletted(state_ids: list[int], volume: int, min_bits: int = 4,
                     global_bits: int = 15, max_local_bits: int = 8) -> bytes:
    """编码一个 PalettedContainer 载荷（与 Java PalettedContainer.read 对齐）。"""
    uniq = sorted(set(state_ids))
    if len(uniq) == 1:
        return bytes([0]) + _varint(uniq[0]) + _varint(0)

    need = (len(uniq) - 1).bit_length()
    if need > max_local_bits:
        # 全局调色板：位宽直接放 state id
        return bytes([global_bits]) + _long_array(_pack(state_ids, global_bits))

    bits = max(min_bits, need)
    idx = {v: i for i, v in enumerate(uniq)}
    out = bytes([bits]) + _varint(len(uniq))
    for v in uniq:
        out += _varint(v)
    out += _long_array(_pack([idx[v] for v in state_ids], bits))
    return out


def _decode_storage(longs: Iterable[int], volume: int, bits: int) -> list[int]:
    """int[] 位数组 → 索引列表（不跨 long）。"""
    longs = list(longs)
    if bits == 0 or not longs:
        return [0] * volume
    vpl = 64 // bits
    mask = (1 << bits) - 1
    out = []
    for i in range(volume):
        li = i // vpl
        if li >= len(longs):
            out.append(0); continue
        off = (i - li * vpl) * bits
        out.append((longs[li] >> off) & mask)
    return out


@dataclass
class SaveChunk:
    """存档里一个区块的、已解析好的方块与生物群系。"""

    cx: int
    cz: int
    status: str = ""
    sections: dict[int, dict] = field(default_factory=dict)   # y_idx -> {"palette": [名字], "states": [4096 名字], "biomes": [64 名字]}
    palette_specs: list[str] = field(default_factory=list)     # 全部去重后的方块名(+属性)规格


def _props_to_spec(name: str, props: dict | None) -> str:
    """{'Name':'minecraft:oak_log','Properties':{'axis':'y'}} → 'minecraft:oak_log[axis=y]'"""
    if not props:
        return name
    parts = ",".join(f"{k}={v}" for k, v in sorted(props.items()))
    return f"{name}[{parts}]"


def parse_save_chunk(cx: int, cz: int, nbt: dict) -> SaveChunk:
    """存档区块 NBT → SaveChunk（方块用**名字规格**表示，state id 留给渲染器解析）。"""
    ch = SaveChunk(cx=cx, cz=cz, status=str(nbt.get("Status", "")))
    specs_seen: set[str] = set()
    by_y: dict[int, dict] = {}

    for s in nbt.get("sections", []) or []:
        sy = s.get("Y")
        if sy is None:
            continue
        bs = s.get("block_states")
        if not isinstance(bs, dict):
            continue
        palette = bs.get("palette") or []
        spec_list = [
            _props_to_spec(str(p.get("Name", "minecraft:air")), p.get("Properties"))
            for p in palette
        ]
        for sp in spec_list:
            specs_seen.add(sp)

        # 存档 data 是「按 long 打包的调色板索引」（无 varint 前缀）
        storage = bs.get("data") or []
        n = len(spec_list)
        if n <= 1:
            states = [spec_list[0] if spec_list else "minecraft:air"] * SECTION_VOLUME
        else:
            bits = max(4, (n - 1).bit_length())
            idxs = _decode_storage(storage, SECTION_VOLUME, bits)
            states = [spec_list[i] if i < n else "minecraft:air" for i in idxs]

        # 生物群系：同样结构，4x4x4
        bio = s.get("biomes") or {}
        bio_pal = bio.get("palette") or []
        bio_specs = [
            (p if isinstance(p, str) else str(p.get("Name", "minecraft:plains")))
            for p in bio_pal
        ]
        if not bio_specs:
            biomes = ["minecraft:plains"] * BIOME_VOLUME
        elif len(bio_specs) == 1:
            biomes = [bio_specs[0]] * BIOME_VOLUME
        else:
            bb = max(1, (len(bio_specs) - 1).bit_length())
            bidx = _decode_storage(bio.get("data") or [], BIOME_VOLUME, bb)
            biomes = [bio_specs[i] if i < len(bio_specs) else "minecraft:plains" for i in bidx]

        by_y[int(sy)] = {"palette": spec_list, "states": states, "biomes": biomes}

    ch.sections = by_y
    ch.palette_specs = sorted(specs_seen)
    return ch


class WorldRenderer:
    """世界 → 画面。持有一个常驻 Java 渲染器进程。"""

    def __init__(self, bridge: RenderBridge, instance: Any | None = None):
        self.bridge = bridge
        self.instance = instance
        self._spec_cache: dict[str, int] = {}
        self._loaded: set[tuple[int, int]] = set()
        self.chunks: dict[tuple[int, int], SaveChunk] = {}

    # ---------------- 构造 ----------------

    @classmethod
    def open(cls, instance: Any = None, *, store: str | None = None,
             classpath: str | None = None, start: bool = True,
             java: str = "java") -> "WorldRenderer":
        """从 instance（mcfold.instance.Instance）或 store 目录打开渲染器。

        classpath 缺省取环境变量 MCFOLD_RENDER_CLASSPATH；
        可先跑 mcfold/scaffold/build.sh 生成。
        """
        game_dir = assets = asset_index = None
        if instance is not None:
            root = getattr(instance, "content_root", None) or getattr(instance, "path", None)
            if root:
                root = Path(root)
                assets = str(root / "assets")
                idx = root / "assets" / "indexes"
                if idx.is_dir():
                    names = sorted(p.stem for p in idx.glob("*.json"))
                    asset_index = names[-1] if names else None
            game_dir = str(getattr(instance, "path", root) if root else "")
        if store:
            store_p = Path(store)
            assets = str(store_p / "assets")
            names = sorted(p.stem for p in (store_p / "assets" / "indexes").glob("*.json"))
            asset_index = names[-1] if names else None
            game_dir = game_dir or str(store_p / "instances" / "default")

        bridge = RenderBridge(
            java=java, classpath=classpath,
            game_dir=game_dir, assets=assets, asset_index=asset_index,
        )
        wr = cls(bridge, instance)
        if start:
            bridge.start()
        return wr

    def close(self) -> None:
        self.bridge.stop()

    def __enter__(self) -> "WorldRenderer":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # ---------------- 方块名 → state id ----------------

    def resolve_specs(self, specs: list[str]) -> dict[str, int]:
        """批量把方块名规格解析成 state id（渲染器侧权威注册表）。"""
        todo = [s for s in dict.fromkeys(specs) if s not in self._spec_cache]
        if todo:
            resp = self.bridge._request({"cmd": "resolve", "blocks": todo}, timeout=120.0)
            ids = resp.get("ids", [])
            for spec, sid in zip(todo, ids):
                self._spec_cache[spec] = int(sid)
        return {s: self._spec_cache[s] for s in specs}

    # ---------------- 装载 ----------------

    def load_save(self, save_dir: str | Path, *, center: tuple[int, int] | None = None,
                  radius: int | None = None, dim: str | None = None,
                  only_status: str = "minecraft:full") -> dict:
        """从存档装载区块（零网络、零登录）。

        radius —— 加载半径（区块）。**缺省跟随视距**（这是关键：
        视距设了 12 但只加载 3 圈数据，画面看起来就像视距没生效）。
        传 None 时用 renderDistance；显式传值则两者取大。
        """
        if radius is None:
            radius = max(3, self.render_distance)
        else:
            radius = max(radius, 1)
        save_dir = Path(save_dir)
        level = anvil.read_level_dat(save_dir)
        data = level.get("Data", {})
        spawn = data.get("SpawnX"), data.get("SpawnY"), data.get("SpawnZ")
        if center is None:
            center = ((spawn[0] or 0) >> 4, (spawn[2] or 0) >> 4)

        regs: dict[tuple[int, int], anvil.RegionFile] = {}
        # ★ 先设视野中心：ClientChunkCache 只保留视野内的区块（实测否则注入被丢弃）
        self.set_viewport(center[0], center[1], radius=radius + 2)

        parsed: list[SaveChunk] = []
        skipped = 0
        for cx in range(center[0] - radius, center[0] + radius + 1):
            for cz in range(center[1] - radius, center[1] + radius + 1):
                rx, rz = cx >> 5, cz >> 5
                key = (rx, rz)
                if key not in regs:
                    p = anvil.region_path(save_dir, dim, rx, rz)
                    if not p.is_file():
                        skipped += 1; continue
                    regs[key] = anvil.RegionFile(p)
                raw = regs[key].chunk_nbt(cx, cz)
                if not raw:
                    skipped += 1; continue
                nbt = anvil.nbt_read(raw)
                if only_status and str(nbt.get("Status", "")) != only_status:
                    skipped += 1; continue
                parsed.append(parse_save_chunk(cx, cz, nbt))

        return self.load_chunks(parsed, meta={
            "level_name": data.get("LevelName"),
            "data_version": data.get("DataVersion"),
            "spawn": spawn,
            "center": center,
            "skipped": skipped,
        })

    def load_chunks(self, chunks: list[SaveChunk], meta: dict | None = None) -> dict:
        """把已解析的区块编码并注入渲染器。"""
        all_specs: list[str] = []
        for c in chunks:
            all_specs.extend(c.palette_specs)
        id_map = self.resolve_specs(all_specs)
        unknown = sum(1 for v in id_map.values() if v < 0)

        items: list[tuple[int, int, bytes]] = []
        for c in chunks:
            payload = self._encode_chunk(c, id_map)
            items.append((c.cx, c.cz, payload))
            self.chunks[(c.cx, c.cz)] = c

        resp = self.bridge.feed_chunks(items) if items else {"injected": 0}
        self._loaded.update((c.cx, c.cz) for c in chunks)
        out = {"injected": len(items), "unknownBlocks": unknown,
               "uniqueSpecs": len(id_map)}
        if meta:
            out.update(meta)
        return out

    def _encode_chunk(self, c: SaveChunk, id_map: dict[str, int]) -> bytes:
        """SaveChunk → packet_map_chunk 的 chunkData 载荷。"""
        out = bytearray()
        for sy in range(SECTION_COUNT):
            # ★ 索引 0..23 ↔ 绝对 Y (-4..19)：存档用绝对 Y，载荷用相对索引
            sec = c.sections.get(sy + MIN_SECTION_Y)
            if sec is None:
                # 空 section：空气
                out += _encode_paletted([0] * SECTION_VOLUME, SECTION_VOLUME)
            else:
                ids = [max(0, id_map.get(s, -1)) for s in sec["states"]]
                non_empty = sum(1 for v in ids if v != 0)
                out += (non_empty & 0xFFFF).to_bytes(2, "big")
                out += _encode_paletted(ids, SECTION_VOLUME)
            # biome section
            bios = sec["biomes"] if sec else ["minecraft:plains"] * BIOME_VOLUME
            bio_pal = sorted(set(bios))
            # 生物群系用全局调色板 id 需要 registry —— 简化：统一用 plains 的 id 1
            # （生物群系只影响草/树叶着色，不影响「能不能渲染出世界」）
            out += bytes([0]) + _varint(1) + _varint(0)
        return bytes(out)

    # ---------------- 视角与截图 ----------------

    @property
    def render_distance(self) -> int:
        """当前视距（区块）。默认 8，可用 set_view(render_distance=) 改。"""
        return getattr(self, "_render_distance", 8)

    def set_render_distance(self, rd: int) -> dict:
        """设置视距（区块）。内部同时设 options + serverRenderDistance +
        ClientChunkCache 的 storage 半径 —— 三者缺一则视距不生效（实测）。"""
        self._render_distance = int(rd)
        return self.bridge.set_view(render_distance=int(rd))

    def set_viewport(self, cx: int, cz: int, radius: int | None = None) -> dict:
        """设定视野中心 —— 必须在注入前设好，否则区块会被 ClientChunkCache 丢弃。"""
        req: dict[str, Any] = {"cmd": "viewport", "cx": cx, "cz": cz}
        if radius is not None:
            req["radius"] = radius
        return self.bridge._request(req)

    def debug(self) -> dict:
        """诊断：玩家坐标/区块/网格编译/视野中心的实际状态。"""
        return self.bridge._request({"cmd": "debug"})

    def camera(self, x: float, y: float, z: float,
               yaw: float = 0.0, pitch: float = 0.0) -> None:
        self.bridge.set_camera(x, y, z, yaw, pitch)

    def camera_above(self, cx: float = 0.0, cz: float = 0.0, height: float = 90.0,
                     yaw: float = 0.0, pitch: float = 70.0) -> None:
        """俯视整片区域（默认视角）。"""
        self.bridge.set_camera(cx, height, cz, yaw, pitch)

    def screenshot(self, out: str | None = None, frames: int = 60) -> RenderResult:
        return self.bridge.render(out=out, frames=frames)