# Sparse block placement, authoritative states and deterministic Litematic transport.
from __future__ import annotations

from functools import lru_cache
import json
import os
from pathlib import Path
import random
import sys

MCIO_ROOT = Path(os.environ.get('MCIO_ROOT', '/work/generator/MCIO'))
STATE_FILE = Path(os.environ.get('MINECRAFT_STATES', '/work/generator/MCRender/resources/1.21.1-blocks.json'))
sys.path.insert(0, str(MCIO_ROOT))
from mcio.sketch import LitematicCanvas, parse_blockstate
from mcio.schematic import load_schematic
from mcio.nbt.nbt import File as NBTFile
from mcio.nbt.tag import Compound, Int, String, Byte, Short, List as NBTList

AIR = 'minecraft:air'


@lru_cache(maxsize=1)
def registry():
    """Load vanilla 1.21.1 property domains and defaults from the installed authority."""
    return json.loads(STATE_FILE.read_text())


@lru_cache(maxsize=None)
def split_state(state):
    """Parse a blockstate while retaining its explicit property set."""
    parsed = parse_blockstate(state).state.unpack()
    return parsed['Name'], parsed.get('Properties', {})


def blockstates_equivalent(left: str, right: str) -> bool:
    """Compare names and explicit properties, ignoring property insertion order."""
    return parse_blockstate(left).state.unpack() == parse_blockstate(right).state.unpack()


@lru_cache(maxsize=None)
def make_state(block, pairs=()):
    """Fill generated states with authoritative defaults and reject invalid values."""
    if '[' in block:
        name, props = split_state(block)
        return make_state(name, tuple(sorted(props.items())))
    short = block.removeprefix('minecraft:')
    if short not in registry():
        raise ValueError(f'unknown Java 1.21.1 block {block}')
    domains, defaults = registry()[short]
    props = dict(defaults)
    for key, value in pairs:
        if key not in domains or str(value) not in domains[key]:
            raise ValueError(f'invalid state {short}.{key}={value!r}')
        props[key] = str(value)
    return 'minecraft:' + short + ('[' + ','.join(f'{k}={v}' for k, v in sorted(props.items())) + ']' if props else '')


def name(state):
    """Read a short block name without parsing all properties."""
    return state.split('[', 1)[0].removeprefix('minecraft:')


def is_air(state):
    return name(state) in ('air', 'cave_air', 'void_air')


def passable(state):
    """Conservative player clearance for the vocabulary emitted by this library."""
    n = name(state)
    return (is_air(state) or n.endswith('_door') and not n.endswith('iron_door') or n.endswith('_carpet') or n in ('ladder', 'short_grass', 'fern', 'snow') or n.endswith('_fence_gate') and split_state(state)[1].get('open') == 'true')


def full_support(state):
    """Identify blocks with a usable full top, excluding decorative partial blocks."""
    n = name(state)
    if is_air(state) or n in ('water', 'lava', 'ladder', 'short_grass', 'fern', 'lily_pad', 'seagrass'):
        return False
    if any(n.endswith(suffix) for suffix in ('_door', '_trapdoor', '_fence', '_fence_gate', '_wall', '_pane', '_carpet', '_bed', '_leaves', '_stairs')):
        return False
    if n.endswith('_slab'):
        return split_state(state)[1].get('type') in ('top', 'double')
    if n.startswith('potted_') or n in ('lantern', 'chain', 'flower_pot', 'torch', 'wall_torch', 'brewing_stand', 'wheat', 'poppy', 'cornflower', 'allium', 'oxeye_daisy', 'dandelion', 'sugar_cane') or 'candle' in n:
        return False
    return True


def floor_support(state):
    return full_support(state) or name(state).endswith('_stairs') and split_state(state)[1].get('half', 'bottom') == 'bottom'


def seal(state):
    """Require full backing for partial roof blocks; closed doors are weather seals."""
    n = name(state)
    return full_support(state) or n.endswith('_door') and split_state(state)[1].get('open') == 'false'


class Grid:
    """Store world coordinates sparsely and derive export bounds from final geometry."""

    def __init__(self, seed=0):
        self.blocks = {}
        self.rng = random.Random(seed)
        self.reservations = {}
        self.metadata = {'schema': 1, 'rooms': [], 'routes': [], 'openings': [], 'supports': [], 'furniture': [], 'lights': [], 'interior': [], 'features': [], 'stairs': []}
        self.inventories = {}

    def get(self, x, y, z):
        return self.blocks.get((x, y, z), AIR)

    def put(self, x, y, z, block, **props):
        if any(type(v) is not int for v in (x, y, z)):
            raise ValueError(f'integer coordinates required: {(x, y, z)}')
        state = make_state(block, tuple(sorted(props.items())))
        if is_air(state):
            self.blocks.pop((x, y, z), None)
            self.inventories.pop((x, y, z), None)
        else:
            self.blocks[x, y, z] = state

    def box(self, bounds, block, **props):
        x0, y0, z0, x1, y1, z1 = bounds
        if x1 < x0 or y1 < y0 or z1 < z0:
            return
        for x in range(x0, x1 + 1):
            for y in range(y0, y1 + 1):
                for z in range(z0, z1 + 1):
                    self.put(x, y, z, block, **props)

    def texture(self, bounds, palette):
        x0, y0, z0, x1, y1, z1 = bounds
        for x in range(x0, x1 + 1):
            for y in range(y0, y1 + 1):
                for z in range(z0, z1 + 1):
                    self.put(x, y, z, self.rng.choice(palette))

    def reserve(self, cells, owner, purpose='passage'):
        for p in cells:
            self.reservations.setdefault(tuple(p), []).append((owner, purpose))

    def route(self, points, owner, headroom=2):
        points = [tuple(p) for p in points]
        self.metadata['routes'].append({'owner': owner, 'points': points, 'headroom': headroom})
        self.reserve([(x, y + dy, z) for x, y, z in points for dy in range(headroom)], owner)

    def free(self, cells):
        return all(tuple(p) not in self.reservations and is_air(self.get(*p)) for p in cells)

    @property
    def bounds(self):
        if not self.blocks:
            return (0, 0, 0, 0, 0, 0)
        axes = list(zip(*self.blocks))
        return tuple(min(a) for a in axes) + tuple(max(a) for a in axes)

    def stair(self, x, y, z, material='spruce', facing='south', half='bottom'):
        self.put(x, y, z, material + '_stairs', facing=facing, half=half, shape='straight')

    def door(self, x, floor, z, side, owner, wood='spruce'):
        for dy, half in ((1, 'lower'), (2, 'upper')):
            self.put(x, floor + dy, z, wood + '_door', facing=side, half=half, hinge='left', open='false')
        self.metadata['openings'].append({'kind': 'door', 'owner': owner, 'cells': [(x, floor + 1, z), (x, floor + 2, z)], 'side': side})

    def lantern(self, x, y, z, owner, hanging=False):
        self.put(x, y, z, 'lantern', hanging=str(hanging).lower())
        self.metadata['lights'].append({'owner': owner, 'point': (x, y, z)})

    def inventory(self, p, items):
        self.inventories[tuple(p)] = items

    def export(self, path: Path, title='Hearthwright'):
        """Write one region and an explicit world-coordinate design sidecar.

        Args:
            path: Destination .litematic path.
            title: Schematic metadata title.

        Returns:
            Path: Written file.
        """
        path = Path(path)
        lo_x, lo_y, lo_z, hi_x, hi_y, hi_z = self.bounds
        origin = (lo_x - 1, min(0, lo_y), lo_z - 1)
        sx, sy, sz = hi_x - origin[0] + 2, hi_y - origin[1] + 2, hi_z - origin[2] + 2
        canvas = LitematicCanvas((sy, sz, sx))
        tiles = []
        for (x, y, z), state in sorted(self.blocks.items()):
            px, py, pz = x - origin[0], y - origin[1], z - origin[2]
            canvas.block((py, pz, px), parse_blockstate(state))
            n = name(state)
            kind = 'bed' if n.endswith('_bed') else n
            if kind in ('bed', 'barrel', 'chest', 'smoker', 'furnace'):
                tag = Compound({'id': String('minecraft:' + kind), 'x': Int(px), 'y': Int(py), 'z': Int(pz)})
                if kind != 'bed':
                    tag['Items'] = NBTList[Compound]([Compound({'Slot': Byte(i), 'id': String('minecraft:' + item), 'count': Int(count)}) for i, (item, count) in enumerate(self.inventories.get((x, y, z), []))])
                if kind in ('smoker', 'furnace'):
                    tag.update({'BurnTime': Short(0), 'CookTime': Short(0), 'CookTimeTotal': Short(200), 'RecipesUsed': Compound()})
                tiles.append(tag)
        schematic = canvas.to_litematica(name=title, author='Hearthwright', description='Composed architecture; design and validation metadata in adjacent JSON.', minecraft_data_version=3955)
        schematic.metadata.time_created = 0
        schematic.metadata.time_modified = 0
        nbt = schematic.write_to_nbt()
        region = next(iter(nbt['Regions'].values()))
        region['TileEntities'] = NBTList[Compound](tiles)
        path.parent.mkdir(parents=True, exist_ok=True)
        NBTFile(nbt).save(path)
        self.metadata['origin'] = origin
        self.metadata['size_xyz'] = (sx, sy, sz)
        self.metadata['inventories'] = [{'point': p, 'items': items} for p, items in sorted(self.inventories.items())]
        path.with_suffix('.json').write_text(json.dumps(self.metadata, indent=2, sort_keys=True) + '\n')
        return path

    def verify(self, path: Path):
        """Compare the entire reloaded non-air volume semantically, including extra cells."""
        other = Grid.load(path)
        if set(self.blocks) != set(other.blocks):
            raise AssertionError('non-air coordinate set changed on reload')
        for p, state in self.blocks.items():
            if not blockstates_equivalent(state, other.blocks[p]):
                raise AssertionError(f'cell {p} changed on reload: {state!r} -> {other.blocks[p]!r}')
        return {'blocks': len(self.blocks), 'states': len(set(self.blocks.values())), 'size_xyz': other.metadata['size_xyz']}

    @classmethod
    def load(cls, path: Path, metadata: dict | Path | None = None):
        """Reload geometry with explicit metadata, defaulting to the adjacent sidecar."""
        path = Path(path)
        if metadata is None:
            metadata = path.with_suffix('.json')
        if isinstance(metadata, (str, Path)):
            metadata = json.loads(Path(metadata).read_text())
        grid = cls()
        grid.metadata = metadata
        ox, oy, oz = metadata['origin']
        schematic = load_schematic(path)
        array = schematic.read_flat(0, schematic.volume).reshape(schematic.size_yzx)
        for index, state in enumerate(schematic.palette):
            if is_air(state):
                continue
            import numpy as np
            for y, z, x in zip(*np.where(array == index)):
                grid.blocks[int(x) + ox, int(y) + oy, int(z) + oz] = state
        for route in metadata['routes']:
            grid.reserve([(x, y + dy, z) for x, y, z in route['points'] for dy in range(route['headroom'])], route['owner'])
        nbt = NBTFile.load_regardless_of_gzipped(path)
        region = next(iter(nbt['Regions'].values()))
        tiles = region.get('TileEntities')
        for tag in tiles.unpack() if tiles is not None else []:
            if 'Items' in tag:
                point = tag['x'] + ox, tag['y'] + oy, tag['z'] + oz
                grid.inventories[point] = [(item['id'].removeprefix('minecraft:'), item['count']) for item in sorted(tag['Items'], key=lambda item: item['Slot'])]
        return grid