#!/usr/bin/env python3
"""Score a candidate Sponge schematic against a ground-truth building.

The primary score is exact block-state intersection-over-union (IoU) over
non-air blocks. Schematic metadata, data version, dimensions, offsets, air
padding, entities, and block-entity payloads are not scored.

Usage:
    python score_schematics.py output.schem answer.schem
"""

from __future__ import annotations

import argparse
import gzip
import io
import json
import struct
from collections import Counter, defaultdict
from pathlib import Path

MAX_ALIGNMENT_PAIRS = 5_000_000
MAX_ALIGNMENT_CANDIDATES = 4096

TAG_END = 0
TAG_BYTE = 1
TAG_SHORT = 2
TAG_INT = 3
TAG_LONG = 4
TAG_FLOAT = 5
TAG_DOUBLE = 6
TAG_BYTE_ARRAY = 7
TAG_STRING = 8
TAG_LIST = 9
TAG_COMPOUND = 10
TAG_INT_ARRAY = 11
TAG_LONG_ARRAY = 12

AIR_STATES = {"minecraft:air", "minecraft:cave_air", "minecraft:void_air"}

Position = tuple[int, int, int]
BlockMap = dict[Position, str]


def read_exact(stream: io.BytesIO, size: int) -> bytes:
    data = stream.read(size)
    if len(data) != size:
        raise EOFError("Unexpected end of NBT stream")
    return data


def read_utf(stream: io.BytesIO) -> str:
    length = struct.unpack(">H", read_exact(stream, 2))[0]
    return read_exact(stream, length).decode("utf-8")


def read_payload(stream: io.BytesIO, tag_type: int) -> object:
    if tag_type == TAG_BYTE:
        return struct.unpack(">b", read_exact(stream, 1))[0]
    if tag_type == TAG_SHORT:
        return struct.unpack(">h", read_exact(stream, 2))[0]
    if tag_type == TAG_INT:
        return struct.unpack(">i", read_exact(stream, 4))[0]
    if tag_type == TAG_LONG:
        return struct.unpack(">q", read_exact(stream, 8))[0]
    if tag_type == TAG_FLOAT:
        return struct.unpack(">f", read_exact(stream, 4))[0]
    if tag_type == TAG_DOUBLE:
        return struct.unpack(">d", read_exact(stream, 8))[0]
    if tag_type == TAG_BYTE_ARRAY:
        length = struct.unpack(">i", read_exact(stream, 4))[0]
        return read_exact(stream, length)
    if tag_type == TAG_STRING:
        return read_utf(stream)
    if tag_type == TAG_LIST:
        element_type = read_exact(stream, 1)[0]
        length = struct.unpack(">i", read_exact(stream, 4))[0]
        if length < 0:
            raise ValueError("NBT list has a negative length")
        return element_type, [
            read_payload(stream, element_type) for _ in range(length)
        ]
    if tag_type == TAG_COMPOUND:
        values: list[tuple[str, int, object]] = []
        while True:
            child_type = read_exact(stream, 1)[0]
            if child_type == TAG_END:
                return values
            name = read_utf(stream)
            values.append((name, child_type, read_payload(stream, child_type)))
    if tag_type == TAG_INT_ARRAY:
        length = struct.unpack(">i", read_exact(stream, 4))[0]
        if length < 0:
            raise ValueError("NBT int array has a negative length")
        return [
            struct.unpack(">i", read_exact(stream, 4))[0]
            for _ in range(length)
        ]
    if tag_type == TAG_LONG_ARRAY:
        length = struct.unpack(">i", read_exact(stream, 4))[0]
        if length < 0:
            raise ValueError("NBT long array has a negative length")
        return [
            struct.unpack(">q", read_exact(stream, 8))[0]
            for _ in range(length)
        ]
    raise ValueError(f"Unsupported NBT tag type {tag_type}")


def read_nbt(raw: bytes) -> list[tuple[str, int, object]]:
    stream = io.BytesIO(raw)
    root_type = read_exact(stream, 1)[0]
    if root_type != TAG_COMPOUND:
        raise ValueError("Schematic NBT root is not a compound")
    read_utf(stream)
    root = read_payload(stream, TAG_COMPOUND)
    if stream.read(1):
        raise ValueError("Unexpected trailing data after NBT root")
    return root  # type: ignore[return-value]


def as_dict(
    compound: list[tuple[str, int, object]],
) -> dict[str, tuple[int, object]]:
    return {name: (tag_type, value) for name, tag_type, value in compound}


def compound_value(
    compound: dict[str, tuple[int, object]],
    name: str,
) -> list[tuple[str, int, object]]:
    try:
        tag_type, value = compound[name]
    except KeyError as error:
        raise ValueError(f"Missing required NBT compound {name!r}") from error
    if tag_type != TAG_COMPOUND:
        raise ValueError(f"NBT tag {name!r} is not a compound")
    return value  # type: ignore[return-value]


def canonical_state(state: str) -> str:
    """Normalize namespace and property ordering without changing semantics."""
    block_id, separator, properties = state.partition("[")
    if ":" not in block_id:
        block_id = f"minecraft:{block_id}"
    if not separator:
        return block_id
    if not properties.endswith("]"):
        raise ValueError(f"Malformed block state: {state}")
    pairs = []
    for property_text in properties[:-1].split(","):
        if "=" not in property_text:
            raise ValueError(f"Malformed block-state property: {state}")
        key, value = property_text.split("=", 1)
        pairs.append((key, value))
    rendered = ",".join(f"{key}={value}" for key, value in sorted(pairs))
    return f"{block_id}[{rendered}]"


def decode_varints(data: bytes) -> list[int]:
    values = []
    value = 0
    shift = 0
    for byte in data:
        value |= (byte & 0x7F) << shift
        if byte & 0x80:
            shift += 7
            if shift > 35:
                raise ValueError("Block-data varint is too large")
        else:
            values.append(value)
            value = 0
            shift = 0
    if shift:
        raise ValueError("Truncated block-data varint")
    return values


def load_schematic(path: Path) -> BlockMap:
    try:
        raw = gzip.decompress(path.read_bytes())
    except OSError as error:
        raise ValueError(f"{path} is not a gzip-compressed schematic") from error

    root = as_dict(read_nbt(raw))
    schematic = (
        as_dict(compound_value(root, "Schematic"))
        if "Schematic" in root
        else root
    )

    width = int(schematic["Width"][1])
    height = int(schematic["Height"][1])
    length = int(schematic["Length"][1])
    if width <= 0 or height <= 0 or length <= 0:
        raise ValueError(f"{path} has invalid dimensions")

    if "Blocks" in schematic:
        block_container = as_dict(compound_value(schematic, "Blocks"))
        palette_compound = as_dict(compound_value(block_container, "Palette"))
        data_type, encoded_data = block_container["Data"]
    else:
        # Sponge schematic v1/v2 layout.
        block_container = schematic
        palette_compound = as_dict(compound_value(schematic, "Palette"))
        data_type, encoded_data = schematic["BlockData"]

    if data_type != TAG_BYTE_ARRAY:
        raise ValueError(f"{path} block data is not an NBT byte array")

    palette_by_id = {
        int(palette_entry[1]): canonical_state(state)
        for state, palette_entry in palette_compound.items()
    }
    palette_ids = decode_varints(bytes(encoded_data))
    expected_entries = width * height * length
    if len(palette_ids) != expected_entries:
        raise ValueError(
            f"{path} contains {len(palette_ids)} blocks; expected {expected_entries}"
        )

    blocks: BlockMap = {}
    layer_size = width * length
    for index, palette_id in enumerate(palette_ids):
        try:
            state = palette_by_id[palette_id]
        except KeyError as error:
            raise ValueError(f"{path} references unknown palette id {palette_id}") from error
        if state in AIR_STATES:
            continue
        y, within_layer = divmod(index, layer_size)
        z, x = divmod(within_layer, width)
        blocks[x, y, z] = state
    return blocks


def minimum_position(blocks: BlockMap) -> Position:
    return (
        min(position[0] for position in blocks),
        min(position[1] for position in blocks),
        min(position[2] for position in blocks),
    )


def add_position(left: Position, right: Position) -> Position:
    return left[0] + right[0], left[1] + right[1], left[2] + right[2]


def subtract_position(left: Position, right: Position) -> Position:
    return left[0] - right[0], left[1] - right[1], left[2] - right[2]


def alignment_candidates(candidate: BlockMap, ground_truth: BlockMap) -> list[Position]:
    """Find likely translations by voting between blocks with equal states."""
    candidate_by_state: dict[str, list[Position]] = defaultdict(list)
    truth_by_state: dict[str, list[Position]] = defaultdict(list)
    for position, state in candidate.items():
        candidate_by_state[state].append(position)
    for position, state in ground_truth.items():
        truth_by_state[state].append(position)

    shared_states = sorted(
        candidate_by_state.keys() & truth_by_state.keys(),
        key=lambda state: (
            len(candidate_by_state[state]) * len(truth_by_state[state]),
            state,
        ),
    )
    votes: Counter[Position] = Counter()
    pairs_seen = 0
    for state in shared_states:
        state_pairs = len(candidate_by_state[state]) * len(truth_by_state[state])
        if pairs_seen and pairs_seen + state_pairs > MAX_ALIGNMENT_PAIRS:
            break
        for candidate_position in candidate_by_state[state]:
            for truth_position in truth_by_state[state]:
                votes[subtract_position(truth_position, candidate_position)] += 1
        pairs_seen += state_pairs

    min_delta = subtract_position(
        minimum_position(ground_truth),
        minimum_position(candidate),
    )
    candidates = [delta for delta, _ in votes.most_common(MAX_ALIGNMENT_CANDIDATES)]
    if min_delta not in candidates:
        candidates.append(min_delta)
    return candidates or [(0, 0, 0)]


def translate(blocks: BlockMap, delta: Position) -> BlockMap:
    return {add_position(position, delta): state for position, state in blocks.items()}


def exact_matches(candidate: BlockMap, ground_truth: BlockMap) -> int:
    return sum(
        ground_truth.get(position) == state
        for position, state in candidate.items()
    )


def best_alignment(candidate: BlockMap, ground_truth: BlockMap) -> tuple[Position, BlockMap]:
    best_delta = (0, 0, 0)
    best_blocks = candidate
    best_key = (-1, -1, 0, 0, 0)
    truth_positions = set(ground_truth)

    for delta in alignment_candidates(candidate, ground_truth):
        aligned = translate(candidate, delta)
        key = (
            exact_matches(aligned, ground_truth),
            len(set(aligned) & truth_positions),
            -abs(delta[0]) - abs(delta[1]) - abs(delta[2]),
            -abs(delta[1]),
            -abs(delta[0]) - abs(delta[2]),
        )
        if key > best_key:
            best_key = key
            best_delta = delta
            best_blocks = aligned
    return best_delta, best_blocks


def block_type(state: str) -> str:
    return state.partition("[")[0]


def set_iou(left: set[object], right: set[object]) -> float:
    union = left | right
    return 1.0 if not union else len(left & right) / len(union)


def score(candidate: BlockMap, ground_truth: BlockMap) -> dict[str, object]:
    if not candidate and not ground_truth:
        delta, aligned = (0, 0, 0), candidate
    elif not candidate or not ground_truth:
        delta, aligned = (0, 0, 0), candidate
    else:
        delta, aligned = best_alignment(candidate, ground_truth)

    candidate_states = {
        (position[0], position[1], position[2], state)
        for position, state in aligned.items()
    }
    truth_states = {
        (position[0], position[1], position[2], state)
        for position, state in ground_truth.items()
    }
    candidate_types = {
        (position[0], position[1], position[2], block_type(state))
        for position, state in aligned.items()
    }
    truth_types = {
        (position[0], position[1], position[2], block_type(state))
        for position, state in ground_truth.items()
    }
    candidate_positions = set(aligned)
    truth_positions = set(ground_truth)

    exact = len(candidate_states & truth_states)
    wrong_state_positions = sum(
        position in aligned
        and position in ground_truth
        and aligned[position] != ground_truth[position]
        for position in candidate_positions & truth_positions
    )
    missing_positions = truth_positions - candidate_positions
    extra_positions = candidate_positions - truth_positions

    exact_state_iou = set_iou(candidate_states, truth_states)
    return {
        "score": exact_state_iou,
        "perfect_match": exact_state_iou == 1.0,
        "alignment": {
            "translation_candidate_to_ground_truth": list(delta),
            "rotation_degrees": 0,
            "mirrored": False,
        },
        "metrics": {
            "exact_block_state_iou": exact_state_iou,
            "block_type_iou": set_iou(candidate_types, truth_types),
            "occupancy_iou": set_iou(candidate_positions, truth_positions),
        },
        "counts": {
            "ground_truth_non_air_blocks": len(ground_truth),
            "candidate_non_air_blocks": len(candidate),
            "exact_block_state_matches": exact,
            "wrong_block_state_positions": wrong_state_positions,
            "missing_positions": len(missing_positions),
            "extra_positions": len(extra_positions),
        },
        "policy": {
            "air_scored": False,
            "metadata_scored": False,
            "dimensions_scored": False,
            "offset_scored": False,
            "entities_scored": False,
            "block_entity_payloads_scored": False,
            "translation_normalized": True,
            "rotation_normalized": False,
            "mirror_normalized": False,
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compare output.schem with answer.schem using strict block-state IoU."
    )
    parser.add_argument(
        "candidate",
        nargs="?",
        type=Path,
        default=Path("output.schem"),
    )
    parser.add_argument(
        "ground_truth",
        nargs="?",
        type=Path,
        default=Path("answer.schem"),
    )
    parser.add_argument(
        "--require-perfect",
        action="store_true",
        help="Exit with status 1 unless the score is exactly 1.",
    )
    parser.add_argument(
        "--compact",
        action="store_true",
        help="Emit compact JSON instead of indented JSON.",
    )
    args = parser.parse_args()

    candidate = load_schematic(args.candidate)
    ground_truth = load_schematic(args.ground_truth)
    report = score(candidate, ground_truth)
    print(
        json.dumps(
            report,
            indent=None if args.compact else 2,
            sort_keys=True,
        )
    )
    if args.require_perfect and not report["perfect_match"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()