# One-region Litematica transport paired with verified structural records.
import hashlib
import json
import sys
from dataclasses import fields, is_dataclass
from pathlib import Path
from . import __version__
from .blocks import canonical, AIR
from .space import Box, Frame
from .randomness import Scope
from .kernel import contracts
from .kernel.scene import Scene, Node, Cell, Limits
from .kernel.errors import ContractError

FORMAT_VERSION = 1
TYPES = {c.__name__: c for c in (Box, Frame, Scope, Node, Cell, Limits, contracts.Contract, contracts.Rule, contracts.Port, contracts.Grant, contracts.Binding, contracts.Capability, contracts.Domain)}


def encode(value):
    if is_dataclass(value):
        return {'$type': type(value).__name__, **{f.name: encode(getattr(value, f.name)) for f in fields(value)}}
    if isinstance(value, tuple):
        return {'$tuple': [encode(v) for v in value]}
    if isinstance(value, set):
        return {'$set': [encode(v) for v in sorted(value)]}
    if isinstance(value, dict):
        return {k: encode(v) for k, v in value.items()}
    if isinstance(value, list):
        return [encode(v) for v in value]
    return value


def decode(value):
    if isinstance(value, list):
        return [decode(v) for v in value]
    if isinstance(value, dict):
        if '$tuple' in value:
            return tuple(decode(v) for v in value['$tuple'])
        if '$set' in value:
            return set(decode(v) for v in value['$set'])
        if '$type' in value:
            return TYPES[value['$type']](**{k: decode(v) for k, v in value.items() if k != '$type'})
        return {k: decode(v) for k, v in value.items()}
    return value


def semantic_digest(scene):
    digest = hashlib.sha256()
    for p, c in sorted(scene.blocks.items()):
        digest.update(json.dumps([p, canonical(c.state), c.nbt], sort_keys=True, separators=(',', ':')).encode())
    return digest.hexdigest()


def companion_path(path: Path):
    return path.with_suffix(path.suffix + '.structure.json')


def _mcio():
    try:
        import mcio
    except ImportError:
        sys.path.insert(0, '/work/generator/MCIO')
    from mcio.sketch import LitematicCanvas
    from mcio.schematic import load_schematic
    from mcio.nbt import tag
    from mcio.nbt.nbt import File
    return LitematicCanvas, load_schematic, tag, File


def _to_nbt(value, tag, key=''):
    if isinstance(value, dict):
        return tag.Compound({k: _to_nbt(v, tag, k) for k, v in value.items()})
    if isinstance(value, list):
        if not value:
            return tag.List[tag.Compound]([])
        items = [_to_nbt(v, tag) for v in value]
        return tag.List[type(items[0])](items)
    if isinstance(value, str):
        return tag.String(value)
    if isinstance(value, int):
        return tag.Byte(value) if key == 'Slot' else tag.Int(value)
    raise TypeError(type(value))


def _plain(value):
    if isinstance(value, dict):
        return {str(k): _plain(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_plain(v) for v in value]
    if isinstance(value, str):
        return str(value)
    return int(value)


def export_scene(scene, path: Path, verify=True):
    path = Path(path)
    scene.finalize()
    bounds = Box.enclosing(scene.blocks)
    if max(bounds.size) > scene.limits.dimension:
        raise ContractError('export-bounds', conditions=str(bounds))
    if bounds.volume > scene.limits.export_volume:
        raise ContractError('export-volume', conditions=f'{bounds.volume} > {scene.limits.export_volume}')
    canvas_type, _, tag, nbtfile = _mcio()
    sx, sy, sz = bounds.size
    canvas = canvas_type((sy, sz, sx))
    offset = bounds.lo
    block_entities = []
    for p, cell in sorted(scene.blocks.items()):
        x, y, z = tuple(v - o for v, o in zip(p, offset))
        canvas.block((y, z, x), cell.state)
        if cell.nbt is not None:
            block_entities.append(_to_nbt({**cell.nbt, 'x': x, 'y': y, 'z': z}, tag))
    schematic = canvas.to_litematica(name='Hearth', author='Hearth procedural library', description='Contract-composed settlement; structural identity in companion JSON', minecraft_data_version=3955)
    nbt = schematic.write_to_nbt()
    nbt['Regions']['Hearth']['TileEntities'] = tag.List[tag.Compound](block_entities)
    nbt['Regions']['Hearth']['Entities'] = tag.List[tag.Compound]([])
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + '.tmp')
    nbtfile(nbt).save(temporary)
    owners = sorted(scene.nodes)
    operations = sorted(scene.operations)
    owner_id, op_id = {v: i for i, v in enumerate(owners)}, {v: i for i, v in enumerate(operations)}
    data = {'format': 'hearth.structure', 'version': FORMAT_VERSION, 'implementation': __version__, 'minecraft_data_version': 3955, 'digest': semantic_digest(scene), 'offset': offset, 'size': bounds.size, 'seed': scene.seed, 'style': scene.style, 'preferences': {k: v.expression() for k, v in scene.preferences.items()}, 'domain': encode(scene.domain), 'limits': encode(scene.limits), 'scene_version': scene.version, 'nodes': encode(scene.nodes), 'relations': encode(scene.relations), 'dependencies': encode(scene.dependencies), 'operations': encode(scene.operations), 'erasures': [[list(p), owner] for p, owner in sorted(scene.erasures.items())], 'displaced': {k: [[list(p), encode(c)] for p, c in sorted(v.items())] for k, v in scene.displaced.items()}, 'events': encode(scene.events), 'port_uses': [[*k, v] for k, v in sorted(scene.port_uses.items())], 'owners': owners, 'operation_ids': operations, 'cells': [[*p, owner_id[c.owner], op_id[c.operation], list(c.shared)] for p, c in sorted(scene.blocks.items())]}
    target = companion_path(path)
    temporary_meta = target.with_suffix('.json.tmp')
    temporary_meta.write_text(json.dumps(data, sort_keys=True, separators=(',', ':')))
    temporary.replace(path)
    temporary_meta.replace(target)
    if verify:
        restored = load_scene(path)
        if semantic_digest(restored) != semantic_digest(scene):
            raise ContractError('roundtrip', conditions='Semantic geometry differs')
    return path


def load_scene(path: Path):
    path = Path(path)
    meta = companion_path(path)
    if not meta.exists():
        raise ContractError('untrusted-structure', conditions='Missing companion; geometry cannot reconstruct semantic identity')
    data = json.loads(meta.read_text())
    if data.get('format') != 'hearth.structure' or data.get('version') != FORMAT_VERSION:
        raise ContractError('untrusted-structure', conditions='Unknown companion format/version')
    _, load, _, nbtfile = _mcio()
    loaded = load(path)
    nbt = nbtfile.load_regardless_of_gzipped(filename=path)
    regions = nbt['Regions']
    if len(regions) != 1 or int(nbt['MinecraftDataVersion']) != 3955:
        raise ContractError('transport', conditions='Require one Java 1.21.1 region')
    region = next(iter(regions.values()))
    if list(loaded.size_yzx) != [data['size'][1], data['size'][2], data['size'][0]]:
        raise ContractError('geometry-mismatch', conditions='Export dimensions changed')
    if region['Entities']:
        raise ContractError('geometry-mismatch', conditions='Unexpected entities')
    inventories = {}
    for entry in region['TileEntities']:
        item = _plain(entry)
        p = tuple(item.pop(k) + data['offset'][i] for i, k in enumerate(('x', 'y', 'z')))
        inventories[p] = item
    from .environment.fields import from_expression
    scene = Scene(data['seed'], decode(data['limits']), decode(data['domain']), data.get('style'), {k: from_expression(v) for k, v in data.get('preferences', {}).items()})
    scene.nodes = decode(data['nodes'])
    scene.relations = decode(data['relations'])
    for child, kind, parent in scene.relations:
        if kind in ('contained_by', 'attached_to'):
            scene.members.setdefault(parent, set()).add(child)
    scene.dependencies = decode(data['dependencies'])
    scene.operations = decode(data['operations'])
    scene.erasures = {tuple(p): owner for p, owner in data.get('erasures', [])}
    scene.events = decode(data['events'])
    scene.displaced = {k: {tuple(p): decode(c) for p, c in v} for k, v in data.get('displaced', {}).items()}
    scene.version = data['scene_version']
    scene.port_uses = {(a, b): n for a, b, n in data['port_uses']}
    ownership = {tuple(row[:3]): row[3:] for row in data['cells']}
    scene.reverse = {k: set() for k in scene.nodes}
    ids = loaded.read_flat(0, loaded.volume).reshape(loaded.size_yzx)
    import numpy as np
    for y, z, x in np.argwhere(ids != 0):
        block = canonical(loaded.palette[int(ids[y, z, x])])
        if block == AIR:
            continue
        p = tuple(int(v) + o for v, o in zip((x, y, z), data['offset']))
        if p not in ownership:
            raise ContractError('geometry-mismatch', position=p, conditions='Unrecorded block')
        owner, op, shared = ownership[p]
        owner = data['owners'][owner]
        scene.blocks[p] = Cell(block, owner, data['operation_ids'][op], inventories.get(p), tuple(shared))
        scene.reverse[owner].add(p)
        for participant in shared:
            scene.reverse_shared.setdefault(participant, set()).add(p)
    if set(ownership) != set(scene.blocks) or set(inventories) - set(scene.blocks) or semantic_digest(scene) != data['digest']:
        raise ContractError('geometry-mismatch', conditions='Schematic and structural companion are not a matching semantic pair')
    scene.finalize()
    return scene