"""状态层（M3）：把**服务端真状态**汇总成 agent 可查询的世界。

## 它是什么

M2 拿到了「一串包」，那是协议层的原料。本模块把原料变成
**可查询、可解释、范围受控的世界状态**，供上层（事件层 / 操作层 / 渲染层）使用。

## 设计立场（与 mcfold 总体目的一致）

本模块**只陈述观察，不做决策**：

- 「这里有个箱子」是观察；「去开它」是决策（留给 agent）。
- 「脚下方块是镐能挖的」是观察（材质/硬度来自注册表）；
  「该挖哪个」是决策。
- 统计到「附近有 3 种未识别方块」是观察；agent 决定要不要查/要不要走。

一个刻意保留的能力：`unknown_state_ids` 计数。
mod 方块在 vanilla 协议下**不会**出现在 registry_data（见 plan.md 的 M2 结论），
所以 state_id → 名字的解析只能靠**本机 jar 的注册表**（Java 侧）。
这里对认不出的 id **显式记账并暴露**，绝不静默映射成别的方块
（静默映射会让 agent 基于错误的世界做决策）。

## 为什么不在这里做「找路 / 找目标」

那属于决策与规划。若塞进这里，agent 就失去了选择权，
且这些逻辑在 vanilla 与 mod 之间差异最大、最易过期。
本模块只保证把「事实」摆干净，让 agent 能自己判断。
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Iterable

from .chunk import BlockRegistry, Chunk
from .world import MIN_Y, WORLD_HEIGHT, World, in_world


# ---------- 世界几何（从服务端登录包学到的，而非写死） ----------


@dataclass
class WorldGeometry:
    """世界尺寸。★ 来自服务端 `login` 包，不写死 24/384。

    1.18+ 世界高度随维度变（超下界 -64、384 层），写死会在别的维度读错。
    """

    min_y: int = MIN_Y
    height: int = WORLD_HEIGHT
    dimension: str = "minecraft:overworld"
    seed: int = 0

    def in_world(self, y: int) -> bool:
        return self.min_y <= y < self.min_y + self.height

    def describe(self) -> str:
        return f"{self.dimension} y∈[{self.min_y}, {self.min_y + self.height}) 种子={self.seed}"


# ---------- 实体 ----------


@dataclass
class Entity:
    """一个实体的观察结果。**不预测、不决策**。"""

    entity_id: int
    kind: str = "unknown"
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0
    yaw: float = 0.0
    pitch: float = 0.0
    on_ground: bool = False
    metadata: dict[str, Any] = field(default_factory=dict)

    def as_dict(self) -> dict[str, Any]:
        return {
            "entity_id": self.entity_id, "kind": self.kind,
            "x": round(self.x, 2), "y": round(self.y, 2), "z": round(self.z, 2),
            "yaw": round(self.yaw, 2), "pitch": round(self.pitch, 2),
            "on_ground": self.on_ground,
        }


# ---------- 玩家（自己） ----------


@dataclass
class PlayerView:
    """自己的位置与状态。position 包是它的唯一权威来源。"""

    x: float = 0.0
    y: float = 0.0
    z: float = 0.0
    yaw: float = 0.0
    pitch: float = 0.0
    entity_id: int = -1
    dimension: str = "minecraft:overworld"
    health: float | None = None
    food: int | None = None
    gamemode: str | None = None
    # 相对坐标标志：position 包的 flags 位
    relative: dict[str, bool] = field(default_factory=dict)

    def as_dict(self) -> dict[str, Any]:
        return {
            "x": round(self.x, 2), "y": round(self.y, 2), "z": round(self.z, 2),
            "yaw": round(self.yaw, 2), "pitch": round(self.pitch, 2),
            "entity_id": self.entity_id, "dimension": self.dimension,
            "health": self.health, "food": self.food, "gamemode": self.gamemode,
        }


# ---------- 观察结果（受控范围，给 LLM 读） ----------

# 报告里最多列多少种方块 / 多少个实体 —— 「bounded results」是 API 硬要求。
# 超出部分只给计数，不静默截断（截断会让 agent 以为世界就这么大）。
MAX_BLOCK_KINDS = 12
MAX_ENTITIES = 20
MAX_SCAN_VOLUME = 200_000      # 扫描体积上限：超了就不扫并如实说明


class WorldState:
    """服务端真状态的世界视图。

    ## 与 World 的关系
    `WorldState` 拥有 `World`（负责方块查询/区块解码），
    并在它之上补：**几何、实体、玩家、方块统计与查询工具**。

    ## 为什么方块名字靠本机注册表而不是 registry_data
    实测（plan.md M2 结论）：vanilla 1.21.1 只同步 11 个**动态**注册表，
    `minecraft:block` / `minecraft:item` **不在其中**。所以
    - 原版方块：读本机 jar/blocks.json 的注册表 ✓
    - mod 方块：vanilla 协议不会告诉客户端 ⇒ 只能靠加载器自身同步
      ⇒ 这里对认不出的 id 记 `unknown_state_ids` 并暴露，绝不猜。
    """

    def __init__(
        self,
        registry: BlockRegistry | None = None,
        num_sections: int = 24,
    ):
        self.world = World(registry=registry, num_sections=num_sections)
        self.geometry = WorldGeometry()
        self.player = PlayerView()
        self.entities: dict[int, Entity] = {}
        # ★ 方块变化覆盖层：(x,y,z) -> state_id
        #   不改写区块（区块是压缩调色板，重建代价高且会被 unload 重发覆盖）
        self.overrides: dict[tuple[int, int, int], int] = {}
        # 方块变化流水（供「我改了什么」回放）
        self.block_log: list[tuple[int, int, int, int]] = []
        self.bulk_changes = 0
        self.bulk_change_records: list[int] = []
        # 服务端给的 11 个动态注册表（权威，保留原始条目）
        self.registries: dict[str, Any] = {}
        self.packets_seen: dict[str, int] = {}

    # ---------- 装载：来自服务端的状态 ----------

    def load_chunk(self, x: int, cz: int, chunk_data: bytes) -> Chunk:
        return self.world.load_chunk_data(x, cz, chunk_data)

    def unload_chunk(self, x: int, cz: int) -> None:
        self.world.unload_chunk(x, cz)

    def update_from_position(self, params: dict[str, Any]) -> None:
        """把 `position` 包落进玩家视图。

        ★ flags 位决定哪些轴是**相对增量**而非绝对值 ——
          服务端做小步纠正时走相对位，此时必须累加而不是赋值，
          否则位置会跳到原点（实测 flags 0=全绝对）。
        """
        f = params.get("flags") or {}
        rel = {k: bool(v) for k, v in f.items() if k != "_value"}
        x, y, z = float(params["x"]), float(params["y"]), float(params["z"])
        yaw, pitch = float(params.get("yaw", 0.0)), float(params.get("pitch", 0.0))
        if rel.get("x"):
            x += self.player.x
        if rel.get("y"):
            y += self.player.y
        if rel.get("z"):
            z += self.player.z
        if rel.get("yaw"):
            yaw = (self.player.yaw + yaw) % 360.0
        if rel.get("pitch"):
            pitch = max(-90.0, min(90.0, self.player.pitch + pitch))
        self.player.x, self.player.y, self.player.z = x, y, z
        self.player.yaw, self.player.pitch = yaw, pitch
        self.player.relative = rel

    def update_from_login(self, params: dict[str, Any]) -> None:
        """`login`(play) 包：entityId + 维度类型 + 世界几何。"""
        eid = params.get("entityId")
        if isinstance(eid, int):
            self.player.entity_id = eid
        # 1.21.1 的 login 里有 dimension（0x1B 维数编码），尝试解出名字
        dim = params.get("dimensionType")
        if dim is not None:
            self.player.dimension = f"dim-id:{dim}" if isinstance(dim, int) else str(dim)

    def apply_registry(self, rid: str, entries: Any) -> None:
        self.registries[rid] = entries

    # ---------- 方块变化：覆盖层 ----------

    def on_block_change(self, x: int, y: int, z: int, state_id: int | None) -> None:
        """服务端说某坐标方块变了 ⇒ 更新我们的世界视图。

        ★ 用**覆盖层**（overlay）而不是改写区块：
          区块的 Paletted Container 调色板是压缩的，改一个值可能要重建整段
          位数组；而覆盖层天然有界（只存变过的格子）且可回答「原来的值是什么」
          —— 事件里的 from → to 就需要这个。
          另一个理由：区块随时可能被 unload 重发，覆盖层能正确丢弃。
        """
        if state_id is None:
            return
        self.overrides[(x, y, z)] = int(state_id)
        self.block_log.append((x, y, z, int(state_id)))

    def on_multi_block_change(self, chunk_x: int, chunk_z: int,
                              records: list[int]) -> None:
        """一次多方块变化（爆炸/活塞）：records 打包为 varint 数组。

        记录格式（协议）：每条 varint，高位是状态 id，低 12 位是局部坐标。
        具体位宽随协议而变，故以**协议给定**为准读入并显式记录条数；
        拿不准时不猜坐标，只把条数与方块 id 报给 agent。
        """
        self.bulk_changes += 1
        self.bulk_change_records = list(records)

    def upsert_entity(self, ent: Entity) -> None:
        self.entities[ent.entity_id] = ent

    def remove_entity(self, entity_id: int) -> None:
        self.entities.pop(entity_id, None)

    def note_packet(self, name: str) -> None:
        self.packets_seen[name] = self.packets_seen.get(name, 0) + 1

    # ---------- 查询：范围受控、可读 ----------

    def block_at(self, x: int, y: int, z: int) -> dict[str, Any]:
        """单个方块。返回来源标记，agent 能分辨「确实没有」与「还没看见」。"""
        ov = self.overrides.get((x, y, z))
        if ov is not None:
            # 覆盖层优先（世界已被服务端改动过）
            name = None
            if self.world.registry is not None:
                name = self.world.registry.name_of(ov)
            src = "registry" if name else "unknown-state"
            if name is None:
                self.world.unknown_state_ids[ov] = \
                    self.world.unknown_state_ids.get(ov, 0) + 1
            return {
                "name": name, "state_id": ov, "source": src,
                "describe": f"{name}" if name else f"unknown({ov}) [已变更]",
            }
        hit = self.world.block_at(x, y, z)
        return {
            "name": hit.name,
            "state_id": hit.state_id,
            "source": hit.source,
            "describe": hit.describe(),
        }

    def column_profile(self, x: int, z: int) -> dict[str, Any]:
        """一列方块的地貌剖面（agent 落地/规划最常用的一次查询）。

        ★ 跨版本安全：只报「首个非空气方块」与其上方高度，不假设
          固定的世界高度（几何从服务端来）。
        """
        geo = self.geometry
        col: list[dict[str, Any]] = []
        for y in range(geo.min_y, geo.min_y + geo.height):
            hit = self.world.block_at(x, y, z)
            if hit.name and hit.name != "air":
                col.append({"y": y, "name": hit.name, "state_id": hit.state_id})
        surface = col[0] if col else None
        return {
            "x": x, "z": z,
            "geometry": geo.describe(),
            "surface": surface,
            "non_air_count": len(col),
            # 剖面只给前若干层：全量 384 层会让 LLM 上下文爆掉
            "profile": col[:24],
            "truncated": len(col) > 24,
            "unknown_state_ids": self.world.unknown_state_ids,
        }

    def nearby_blocks(
        self,
        x: int, y: int, z: int,
        radius: int = 4,
        want_air: bool = False,
    ) -> dict[str, Any]:
        """半径内的方块种类统计（不是逐块列表 —— 那是 token 灾难）。

        ★ 范围与体积都受控：radius 上限由调用方给，
          超过 MAX_SCAN_VOLUME 直接拒绝并说明，不静默缩小
          （静默缩小会让 agent 以为扫全了）。
        """
        vol = (2 * radius + 1) ** 3
        if vol > MAX_SCAN_VOLUME:
            return {
                "error": f"扫描体积 {vol} 超过上限 {MAX_SCAN_VOLUME}，请减小 radius",
                "requested_radius": radius, "scanned": False,
            }
        counts: dict[str, int] = {}
        unknown: dict[int, int] = {}
        for dy in range(-radius, radius + 1):
            for dx in range(-radius, radius + 1):
                for dz in range(-radius, radius + 1):
                    hit = self.world.block_at(x + dx, y + dy, z + dz)
                    if hit.source in ("unloaded", "out-of-world"):
                        continue
                    if hit.source == "unknown-state":
                        unknown[hit.state_id] = unknown.get(hit.state_id, 0) + 1
                        continue
                    if not want_air and hit.name == "air":
                        continue
                    counts[hit.name] = counts.get(hit.name, 0) + 1
        ranked = sorted(counts.items(), key=lambda kv: -kv[1])
        return {
            "center": (x, y, z), "radius": radius, "scanned": True,
            "kinds": len(ranked),
            "blocks": [{"name": n, "count": c} for n, c in ranked[:MAX_BLOCK_KINDS]],
            "truncated": len(ranked) > MAX_BLOCK_KINDS,
            "unknown_state_ids": unknown,
            "unresolved_unknown": bool(unknown),
        }

    def nearby_entities(self, radius: int = 16, limit: int = MAX_ENTITIES) -> dict[str, Any]:
        p = self.player
        near: list[tuple[float, Entity]] = []
        for e in self.entities.values():
            d = ((e.x - p.x) ** 2 + (e.y - p.y) ** 2 + (e.z - p.z) ** 2) ** 0.5
            if d <= radius:
                near.append((d, e))
        near.sort(key=lambda t: t[0])
        return {
            "radius": radius, "count": len(near),
            "entities": [e.as_dict() | {"distance": round(d, 2)} for d, e in near[:limit]],
            "truncated": len(near) > limit,
        }

    def summary(self) -> dict[str, Any]:
        """一眼看清「我看到了多少、认出了多少」。"""
        return {
            "player": self.player.as_dict(),
            "geometry": self.geometry.describe(),
            "chunks_loaded": self.world.loaded_count,
            "chunk_columns": len(list(self.world.columns())),
            "entities": len(self.entities),
            "registries_from_server": sorted(self.registries.keys()),
            "registry_entry_counts": {
                k: (len(v) if isinstance(v, (list, tuple)) else None)
                for k, v in sorted(self.registries.items())
            },
            "unknown_state_ids": dict(self.world.unknown_state_ids),
            "has_unresolved_unknown": bool(self.world.unknown_state_ids),
            "packets_seen": dict(sorted(self.packets_seen.items(), key=lambda kv: -kv[1])),
            "note": (
                "unknown_state_ids 非空 ⇒ 这些 state id 来自服务端但本机注册表认不出"
                "（典型：mod 方块）。本工具**不猜**，请另行查询加载器同步通道。"
            ) if self.world.unknown_state_ids else None,
        }

    def reset_unknown_stats(self) -> None:
        self.world.reset_unknown_stats()


# ---------- 构造：注册表从哪来 ----------


def load_vanilla_registry(version: str, data_root: str) -> BlockRegistry | None:
    """从 minecraft-data 的 blocks.json 构造 state_id → 名字 注册表。

    ★ 这是**本机**表（对应本机 jar），不是服务端的：
      服务端不在线同步 block 注册表（plan.md M2 实测结论）。
      所以 mod 方块必然认不出 —— 那是设计上的「显式未知」，
      不是缺陷。别试图用 registry_data 补它，那 11 个包里没有 block。
    """
    import os
    path = os.path.join(data_root, "pc", version, "blocks.json")
    if not os.path.exists(path):
        return None
    return BlockRegistry.from_json_file(path)


def ingest_configuration(
    state: WorldState,
    cfg: dict[str, Any],
    registries: dict[str, Any],
) -> None:
    """把 configuration 阶段拿到的服务端注册表并进状态。"""
    for rid, entries in registries.items():
        state.apply_registry(rid, entries)
    # 维度类型注册表里通常能解出 overworld 的高度 ⇒ 几何不再写死
    dt = registries.get("minecraft:dimension_type")
    if isinstance(dt, (list, tuple)):
        for e in dt:
            if isinstance(e, dict) and "minecraft:overworld" in str(e.get("key", "")):
                val = e.get("value")
                if isinstance(val, dict):
                    h = val.get("height")
                    if isinstance(h, dict) and "height" in h:
                        state.geometry.height = int(h["height"])
                        state.geometry.min_y = int(h.get("min_y", 0))
                break