# Generate and verify a furnished stone cottage for Minecraft Java 1.21.1.
from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path
from typing import Dict, Iterable, Optional, Tuple

sys.path.insert(0, "/work/generator/MCIO")

from mcio.schematic import load_schematic # noqa: E402
from mcio.sketch import LitematicCanvas, Material, parse_blockstate # noqa: E402
from mcio.nbt.nbt import File as NBTFile # noqa: E402
from mcio.nbt.tag import Byte, Compound, Int, List as NBTList, Short, String # noqa: E402

CONFIG = {
    "seed": 20260923,
    "size_xyz": (16, 15, 14), # width X, height Y, depth Z
    "minecraft_version": "1.21.1",
    "data_version": 3955,
    "output": "output.litematic",
    "name": "Sunflower Stone Cottage",
    "author": "generator",
    "description": "Compact stone and oak cottage, recessed north entry, two side windows, furnished room and ladder loft.",
}

WALL = ["minecraft:andesite"] * 5 + ["minecraft:stone"] * 3 + ["minecraft:cobblestone"] * 2
BASE = ["minecraft:cobblestone"] * 4 + ["minecraft:mossy_cobblestone"] * 2 + ["minecraft:andesite"]
ROOF = ["minecraft:andesite_stairs"] * 6 + ["minecraft:cobblestone_stairs"] * 3 + ["minecraft:stone_stairs"]
TIMBER = "minecraft:stripped_oak_log"
TRIM = "minecraft:spruce_planks"
HOUSE = (3, 3, 11, 11) # x0, z0, x1, z1
WINDOW_Z = range(6, 9)
LOFT_BOUNDS = (5, 6, 8, 9, 6, 10)
LADDER_XZ = (7, 10)
BLOCK_REGISTRY = json.loads(Path("/work/generator/MCRender/resources/1.21.1-blocks.json").read_text())

Position = Tuple[int, int, int]


class Build:
    """Hold a block volume and write it out as a Litematica schematic."""

    def __init__(self, size_xyz: Tuple[int, int, int] = CONFIG["size_xyz"], seed: int = CONFIG["seed"]) -> None:
        """Create an empty volume with one seeded random source.

        Args:
            size_xyz (Tuple[int, int, int]): Volume size as width X, height Y, depth Z.
            seed (int): Seed for every random decision in this build.
        """
        self.size_x, self.size_y, self.size_z = size_xyz
        self.rng = random.Random(seed)
        self.canvas = LitematicCanvas((self.size_y, self.size_z, self.size_x))
        self.placed: Dict[Position, str] = {}

    def inside(self, x: int, y: int, z: int) -> bool:
        """Report whether a coordinate is inside the volume.

        Args:
            x (int): Width coordinate.
            y (int): Height coordinate.
            z (int): Depth coordinate.

        Returns:
            bool: True when every axis is within range.
        """
        return 0 <= x < self.size_x and 0 <= y < self.size_y and 0 <= z < self.size_z

    def put(self, x: int, y: int, z: int, block: str | Material, **properties: str) -> None:
        """Place one block, in XYZ order, with an immediate bounds check.

        Args:
            x (int): Width coordinate.
            y (int): Height coordinate.
            z (int): Depth coordinate.
            block (str | Material): Namespaced block name, a full blockstate string, or a Material.
            **properties: Blockstate properties such as facing or half.

        Raises:
            AssertionError: When the coordinate is outside the volume, naming the coordinate.
        """
        assert self.inside(x, y, z), f"out of bounds: (x={x}, y={y}, z={z}) in {self.size_x}x{self.size_y}x{self.size_z}"
        material = self.material(block, **properties) if isinstance(block, str) else block
        self.canvas.block((y, z, x), material)
        self.placed[(x, y, z)] = material.blockstate

    def material(self, name: str, **properties: str) -> Material:
        """Build a Material from a block name and its properties.

        Args:
            name (str): Namespaced block name, or a full blockstate string.
            **properties: Blockstate properties.

        Returns:
            Material: Material ready for placement.
        """
        if "[" in name:
            return parse_blockstate(name)
        allowed, defaults = BLOCK_REGISTRY[name.removeprefix("minecraft:")]
        complete = dict(defaults, **properties)
        for key, value in complete.items():
            assert key in allowed and value in allowed[key], f"invalid block state: {name} {key}={value}"
        return self.canvas.material(name, **dict(sorted(complete.items())))

    def box(self, bounds: Tuple[int, int, int, int, int, int], block: str | Material, **properties: str) -> None:
        """Fill an inclusive box.

        Args:
            bounds (Tuple[int, int, int, int, int, int]): x0, y0, z0, x1, y1, z1, inclusive on both ends.
            block (str | Material): Block to place.
            **properties: Blockstate properties.
        """
        x0, y0, z0, x1, y1, z1 = bounds
        for x in range(min(x0, x1), max(x0, x1) + 1):
            for y in range(min(y0, y1), max(y0, y1) + 1):
                for z in range(min(z0, z1), max(z0, z1) + 1):
                    self.put(x, y, z, block, **properties)

    def shell(self, bounds: Tuple[int, int, int, int, int, int], block: str | Material, **properties: str) -> None:
        """Fill only the faces of a box, leaving the inside empty.

        Args:
            bounds (Tuple[int, int, int, int, int, int]): x0, y0, z0, x1, y1, z1, inclusive.
            block (str | Material): Block to place.
            **properties: Blockstate properties.
        """
        x0, y0, z0, x1, y1, z1 = bounds
        for x in range(min(x0, x1), max(x0, x1) + 1):
            for y in range(min(y0, y1), max(y0, y1) + 1):
                for z in range(min(z0, z1), max(z0, z1) + 1):
                    on_face = x in (x0, x1) or y in (y0, y1) or z in (z0, z1)
                    if on_face:
                        self.put(x, y, z, block, **properties)

    def texture(self, bounds: Tuple[int, int, int, int, int, int], palette: Iterable[str]) -> None:
        """Fill a box with a weighted mix of related blocks, so the surface is not flat.

        Args:
            bounds (Tuple[int, int, int, int, int, int]): x0, y0, z0, x1, y1, z1, inclusive.
            palette (Iterable[str]): Block names. Repeat a name to make it more common.
        """
        choices = list(palette)
        x0, y0, z0, x1, y1, z1 = bounds
        for x in range(min(x0, x1), max(x0, x1) + 1):
            for y in range(min(y0, y1), max(y0, y1) + 1):
                for z in range(min(z0, z1), max(z0, z1) + 1):
                    self.put(x, y, z, self.rng.choice(choices))

    def get(self, x: int, y: int, z: int) -> Optional[str]:
        """Read back what was placed at a coordinate.

        Args:
            x (int): Width coordinate.
            y (int): Height coordinate.
            z (int): Depth coordinate.

        Returns:
            Optional[str]: Blockstate string, or None when the cell is empty.
        """
        return self.placed.get((x, y, z))

    def build(self) -> "Build":
        """Build the reference cottage and a compact, accessible furnished interior.

        Returns:
            Build: This instance, so calls can be chained.
        """
        self.ground_and_walls()
        self.roof()
        self.openings()
        self.interior()
        self.garden()
        return self

    def ground_and_walls(self) -> None:
        """Place grass, a continuous floor, stone panels and the exposed frame."""
        self.box((0, 0, 0, self.size_x - 1, 0, self.size_z - 1), "minecraft:grass_block")
        self.texture((3, 1, 3, 11, 1, 11), BASE)
        self.texture((4, 1, 4, 10, 1, 10), ["minecraft:oak_planks"] * 4 + [TRIM])
        for x in range(3, 12):
            for z in range(3, 12):
                if x in (3, 11) or z in (3, 11):
                    self.texture((x, 2, z, x, 5, z), WALL)
                    if self.rng.random() < 0.3:
                        self.put(x, 2, z, self.rng.choice(BASE))
        for x in (3, 11):
            for z in (3, 11):
                self.box((x, 1, z, x, 6, z), TIMBER, axis="y")
            self.box((x, 6, 3, x, 6, 11), TIMBER, axis="z")
        for z in (3, 11):
            self.box((4, 6, z, 10, 6, z), TIMBER, axis="x")
        # One small inlaid floor runner keeps the main path level and unobstructed.
        self.box((6, 1, 6, 8, 1, 8), "minecraft:terracotta")
        self.box((7, 1, 6, 7, 1, 8), "minecraft:brown_terracotta")

    def roof(self) -> None:
        """Raise one pale gable with a continuous solid backing and wooden verge."""
        for x in range(3, 12):
            underside = 11 - abs(x - 7)
            for z in (3, 11):
                self.texture((x, 7, z, x, underside, z), WALL)
        for z in (3, 11):
            self.box((7, 7, z, 7, 11, z), TIMBER, axis="y")
            self.box((5, 8, z, 9, 8, z), TIMBER, axis="x")
        for x in (2, 12):
            self.box((x, 6, 2, x, 6, 12), TIMBER, axis="z")
        for x in range(2, 13):
            height = 12 - abs(x - 7)
            for z in range(2, 13):
                verge = z in (2, 12)
                self.put(x, height - 1, z, TRIM if verge else self.rng.choice(WALL))
                if x in (2, 12):
                    self.put(x, height - 1, z, TIMBER, axis="z")
                if x == 7:
                    self.put(x, height, z, "minecraft:spruce_slab" if verge else "minecraft:andesite_slab", type="bottom")
                else:
                    self.put(x, height, z, "minecraft:spruce_stairs" if verge else self.rng.choice(ROOF), facing="east" if x < 7 else "west", half="bottom", shape="straight")
        # Small supported corbels and end caps emphasize the reference's heavy eaves.
        for x in (2, 12):
            for z in (3, 5, 9, 11):
                self.put(x, 5, z, "minecraft:spruce_stairs", facing="west" if x == 2 else "east", half="top", shape="straight")
        self.box((7, 12, 0, 7, 12, 3), TIMBER, axis="z")
        self.put(7, 13, 2, "minecraft:spruce_fence")
        self.put(7, 11, 0, "minecraft:chain", axis="y")
        self.put(7, 10, 0, "minecraft:lantern", hanging="true")
        self.put(7, 12, 12, TIMBER, axis="z")

    def openings(self) -> None:
        """Frame two clear glazed windows and a door recessed behind the facade."""
        for x in (3, 11):
            for z in (5, 9):
                self.box((x, 2, z, x, 5, z), TIMBER, axis="y")
            for y in (2, 5):
                self.box((x, y, 6, x, y, 8), TIMBER, axis="z")
            for z in WINDOW_Z:
                for y in (3, 4):
                    self.put(x, y, z, "minecraft:glass_pane", north="true", south="true", east="false", west="false")
        for x in (6, 8):
            self.box((x, 1, 2, x, 1, 4), "minecraft:stone_bricks")
            self.box((x, 2, 3, x, 3, 4), TIMBER, axis="y")
            self.put(x, 2, 2, "minecraft:oak_fence", north="false", south="true", east="false", west="false")
            self.put(x, 3, 2, "minecraft:oak_fence", north="false", south="true", east="false", west="false")
        self.box((5, 4, 2, 9, 4, 3), TIMBER, axis="x")
        self.box((6, 4, 4, 8, 4, 4), TIMBER, axis="x")
        self.box((7, 2, 2, 7, 3, 3), "minecraft:air")
        self.box((7, 1, 2, 7, 1, 4), "minecraft:stone_bricks")
        self.put(7, 1, 1, "minecraft:stone_brick_stairs", facing="south", half="bottom", shape="straight")
        for half, y in (("lower", 2), ("upper", 3)):
            self.put(7, y, 4, "minecraft:spruce_door", half=half, facing="north", hinge="left", open="false", powered="false")
        self.box((5, 5, 1, 5, 5, 3), TIMBER, axis="z")
        self.put(5, 4, 1, "minecraft:lantern", hanging="true")
        self.put(9, 5, 2, "minecraft:spruce_slab", type="bottom")

    def interior(self) -> None:
        """Furnish one main room and a small rear loft reached by a fixed ladder."""
        self.texture((4, 1, 4, 5, 1, 5), ["minecraft:stone_bricks", "minecraft:andesite"])
        self.put(4, 2, 4, "minecraft:furnace", facing="east", lit="false")
        self.put(4, 2, 5, "minecraft:smoker", facing="east", lit="false")
        self.put(4, 2, 9, "minecraft:crafting_table")
        self.box((4, 4, 4, 4, 4, 5), "minecraft:stone_brick_stairs", facing="west", half="top", shape="straight")
        self.put(4, 3, 4, "minecraft:stone_bricks")
        self.put(4, 2, 10, "minecraft:barrel", facing="east", open="false")
        self.put(4, 3, 10, "minecraft:potted_fern")
        self.put(10, 2, 10, "minecraft:barrel", facing="north", open="false")
        for part, z in (("foot", 9), ("head", 10)):
            self.put(9, 2, z, "minecraft:red_bed", facing="south", part=part, occupied="false")
        self.put(9, 2, 6, "minecraft:spruce_stairs", facing="east", half="top", shape="straight")
        self.put(10, 2, 6, "minecraft:spruce_stairs", facing="west", half="top", shape="straight")
        self.put(9, 2, 5, "minecraft:oak_stairs", facing="north", half="bottom", shape="straight")
        self.put(9, 3, 6, "minecraft:potted_dandelion")
        self.box(LOFT_BOUNDS, TRIM)
        self.box((3, 6, 7, 11, 6, 7), TIMBER, axis="x")
        for x in range(5, 10):
            self.put(x, 7, 8, "minecraft:spruce_fence", east="true" if x < 9 else "false", west="true" if x > 5 else "false", north="false", south="false")
        for y in range(2, 8):
            self.put(7, y, 10, "minecraft:ladder", facing="north")
        self.put(5, 7, 10, "minecraft:barrel", facing="east", open="false")
        self.put(5, 7, 9, "minecraft:bookshelf")
        self.put(9, 7, 10, "minecraft:barrel", facing="north", open="false")
        self.put(7, 5, 7, "minecraft:chain", axis="y")
        self.put(7, 4, 7, "minecraft:lantern", hanging="true")
        self.put(5, 5, 10, "minecraft:chain", axis="y")
        self.put(5, 4, 10, "minecraft:lantern", hanging="true")
        self.put(8, 9, 10, "minecraft:lantern", hanging="true")

    def garden(self) -> None:
        """Add the sunflower planter, supported greenery, hay and a short path."""
        for z in WINDOW_Z:
            self.put(12, 1, z, "minecraft:grass_block")
            self.put(13, 1, z, "minecraft:spruce_trapdoor", facing="east", half="bottom", open="true", powered="false")
        for z, facing in ((5, "north"), (9, "south")):
            self.put(12, 1, z, "minecraft:spruce_trapdoor", facing=facing, half="bottom", open="true", powered="false")
        for z in (6, 7):
            self.put(12, 2, z, "minecraft:cornflower")
        self.put(12, 2, 8, "minecraft:sunflower", half="lower")
        self.put(12, 3, 8, "minecraft:sunflower", half="upper")
        # Leaves touch frame or connected leaves; persistent prevents decay.
        leaves = [(13, 6, 3), (12, 5, 3), (12, 4, 3), (12, 5, 4), (13, 5, 3), (13, 6, 11), (12, 5, 11), (12, 4, 11), (13, 5, 11), (12, 5, 10), (2, 5, 3), (2, 4, 3), (2, 5, 4), (1, 6, 3), (2, 2, 10), (2, 1, 10), (2, 1, 11)]
        for p in leaves:
            self.put(*p, "minecraft:oak_leaves", persistent="true", distance="1")
        for y in (2, 3, 4):
            self.put(12, y, 4, "minecraft:vine", west="true", east="false", north="false", south="false", up="false")
        for y in (2, 3):
            self.put(2, y, 4, "minecraft:vine", east="true", west="false", north="false", south="false", up="false")
        for x, y, z in [(14, 1, 10), (15, 1, 10), (14, 1, 11), (15, 1, 11), (14, 2, 11)]:
            self.put(x, y, z, "minecraft:hay_block", axis="y" if y == 2 else "x")
        path = [(6, 0, 0), (7, 0, 0), (8, 0, 0), (6, 0, 1), (8, 0, 1), (6, 0, 2), (8, 0, 2)]
        for p in path:
            self.put(*p, self.rng.choice(["minecraft:gravel", "minecraft:coarse_dirt", "minecraft:andesite"]))
        tufts = [(1, 1, 1), (3, 1, 1), (10, 1, 1), (12, 1, 2), (14, 1, 4), (14, 1, 7), (15, 1, 8), (13, 1, 12), (10, 1, 13), (8, 1, 13), (3, 1, 13), (1, 1, 8), (1, 1, 5)]
        for p in tufts:
            self.put(*p, "minecraft:short_grass")
        for p in [(2, 1, 6), (4, 1, 12), (15, 1, 5)]:
            self.put(*p, "minecraft:dandelion")

    def export(self, path: Path) -> Path:
        """Write the schematic so that two runs produce identical bytes.

        Args:
            path (Path): Destination .litematic path.

        Returns:
            Path: The path written.
        """
        schematic = self.canvas.to_litematica(
            name=CONFIG["name"],
            author=CONFIG["author"],
            description=CONFIG["description"],
            minecraft_data_version=CONFIG["data_version"],
        )
        # The metadata otherwise records the current wall clock, which changes the bytes on every run.
        schematic.metadata.time_created = 0
        schematic.metadata.time_modified = 0
        path.parent.mkdir(parents=True, exist_ok=True)
        nbt = schematic.write_to_nbt()
        region = next(iter(nbt["Regions"].values()))
        region["TileEntities"] = self.block_entities()
        NBTFile(nbt).save(path)
        return path

    def block_entities(self) -> NBTList:
        """Create bed and container NBT, including a few useful cottage supplies.

        Returns:
            NBTList: Block entities with region-relative coordinates.
        """
        entities = []
        supplies = {(4, 2, 10): ("minecraft:bread", 6), (10, 2, 10): ("minecraft:coal", 12), (5, 7, 10): ("minecraft:wheat_seeds", 16), (9, 7, 10): ("minecraft:oak_log", 8)}
        for (x, y, z), state in sorted(self.placed.items()):
            name = state.split("[")[0].removeprefix("minecraft:")
            kind = "bed" if name.endswith("_bed") else name
            if kind not in ("bed", "furnace", "smoker", "barrel"):
                continue
            tag = Compound({"id": String(f"minecraft:{kind}"), "x": Int(x), "y": Int(y), "z": Int(z)})
            if kind != "bed":
                items = []
                if (x, y, z) in supplies:
                    item_id, count = supplies[(x, y, z)]
                    items.append(Compound({"Slot": Byte(0), "id": String(item_id), "count": Int(count)}))
                tag["Items"] = NBTList[Compound](items)
            if kind in ("furnace", "smoker"):
                tag.update({"BurnTime": Short(0), "CookTime": Short(0), "CookTimeTotal": Short(200), "RecipesUsed": Compound()})
            entities.append(tag)
        return NBTList[Compound](entities)

    def verify(self, path: Path) -> Dict[str, object]:
        """Reload the written file and confirm it matches what was placed.

        Args:
            path (Path): Schematic to reload.

        Returns:
            Dict[str, object]: Size, block count and palette size of the reloaded file.

        Raises:
            AssertionError: When a reloaded cell differs from what was placed.
        """
        loaded = load_schematic(path)
        size_y, size_z, size_x = loaded.size_yzx
        assert (size_x, size_y, size_z) == (self.size_x, self.size_y, self.size_z), "size changed on reload"
        ids = loaded.read_flat(0, loaded.volume).reshape(loaded.size_yzx)
        for (x, y, z), state in self.placed.items():
            assert loaded.palette[int(ids[y, z, x])] == state, f"cell (x={x}, y={y}, z={z}) changed on reload"
        return {
            "size_xyz": [size_x, size_y, size_z],
            "placed_blocks": len(self.placed),
            "palette_states": len(loaded.palette),
        }


def main() -> None:
    """Generate the build, export it, and verify the exported file."""
    parser = argparse.ArgumentParser(description="Generate a Minecraft schematic.")
    parser.add_argument("--output", type=Path, default=Path(CONFIG["output"]))
    parser.add_argument("--seed", type=int, default=CONFIG["seed"])
    args = parser.parse_args()

    build = Build(seed=args.seed).build()
    path = build.export(args.output)
    report = build.verify(path)
    print(f"wrote {path}: {report}")


if __name__ == "__main__":
    main()
