#!/usr/bin/env python3
"""
Generate an MCC BlockPalette C# file from server-generated blocks.json.

The blocks.json file is generated by running:
    java -DbundlerMainClass=net.minecraft.data.Main -jar server.jar --reports

Usage:
    python3 tools/gen_block_palette.py <blocks.json> <suffix>

Example:
    python3 tools/gen_block_palette.py /tmp/mc_reports/reports/blocks.json 1219
    # Generates Palette1219.cs
"""

import json
import math
import re
import sys
from collections.abc import Sequence
from pathlib import Path

OUTPUT_DIR = (Path(__file__).resolve().parent.parent /
              "MinecraftClient" / "Mapping" / "BlockPalettes")
MATERIAL_CS = OUTPUT_DIR.parent / "Material.cs"

# Minecraft renamed these registry keys after the corresponding MCC Material
# names had already stabilized. Keep historical reports mapped to the current
# enum names instead of generating references to members that do not exist.
MATERIAL_NAME_ALIASES = {
    "Grass": "ShortGrass",
    "GrassPath": "DirtPath",
    "Sign": "OakSign",
    "WallSign": "OakWallSign",
}

OUTPUT_FILE_ALIASES = {
    "120": "BlockPalette120.cs",
}


def mc_name_to_csharp(mc_name: str) -> str:
    """Convert minecraft:snake_case to PascalCase C# enum name."""
    name = mc_name.removeprefix("minecraft:")
    csharp_name = "".join(word.capitalize() for word in name.split("_"))
    return MATERIAL_NAME_ALIASES.get(csharp_name, csharp_name)


def load_known_materials() -> set[str]:
    known = set()
    if MATERIAL_CS.exists():
        with MATERIAL_CS.open(encoding="utf-8-sig") as f:
            for line in f:
                m = re.match(r'\s+(\w+),?\s*$', line)
                if m:
                    known.add(m.group(1))
    return known


def get_state_property_definitions(
    block_key: str,
    states: list[dict],
    properties: dict[str, list[str]],
) -> list[tuple[str, list[str], int]]:
    """Build and verify the compact stride schema against every reported state."""
    if not properties:
        return []

    ordered_states = sorted(states, key=lambda state: state["id"])
    first_state_id = ordered_states[0]["id"]
    expected_ids = list(range(first_state_id, first_state_id + len(ordered_states)))
    actual_ids = [state["id"] for state in ordered_states]
    if actual_ids != expected_ids:
        raise ValueError(f"{block_key} has non-contiguous state IDs")

    expected_count = math.prod(len(values) for values in properties.values())
    if expected_count != len(ordered_states):
        raise ValueError(
            f"{block_key} has {len(ordered_states)} states but its properties describe {expected_count} combinations"
        )

    definitions = []
    for name, values in properties.items():
        stride = next(
            (
                candidate
                for candidate in range(1, len(ordered_states) + 1)
                if all(
                    state.get("properties", {}).get(name)
                    == values[(offset // candidate) % len(values)]
                    for offset, state in enumerate(ordered_states)
                )
            ),
            None,
        )
        if stride is None:
            raise ValueError(f"{block_key} property {name} has no regular state stride")
        definitions.append((name, values, stride))

    return definitions


def csharp_string(value: str) -> str:
    """Encode a Python string as a compatible C# string literal."""
    return json.dumps(value, ensure_ascii=False)


def render_state_definitions(
    block_ranges: Sequence[tuple[int, int, str, list[tuple[str, list[str], int]]]],
) -> tuple[list[str], int]:
    """Render the generated state-property section and return its definition count."""
    lines = [
        "        // <auto-generated block-state-properties>",
        "        private static readonly BlockStateDefinition[] stateDefinitions =",
        "        [",
    ]
    property_definition_count = 0
    for min_s, max_s, _, properties in block_ranges:
        if not properties:
            continue

        property_definition_count += 1
        lines.append(f"            new({min_s}, {max_s - min_s + 1},")
        lines.append("            [")
        for index, (name, values, stride) in enumerate(properties):
            encoded_values = ", ".join(csharp_string(value) for value in values)
            suffix = "," if index < len(properties) - 1 else ""
            lines.append(f"                new({csharp_string(name)}, [{encoded_values}], {stride}){suffix}")
        lines.append("            ]),")

    lines += [
        "        ];",
        "        // </auto-generated block-state-properties>",
    ]
    return lines, property_definition_count


def update_existing_palette(output_path: Path, state_lines: list[str]) -> bool:
    """Replace only generated state metadata while preserving established material mappings."""
    if not output_path.exists():
        return False

    source = output_path.read_text(encoding="utf-8")
    dictionary_method = "        protected override Dictionary<int, Material> GetDict()"
    dictionary_index = source.find(dictionary_method)
    if dictionary_index < 0:
        raise ValueError(f"{output_path} does not contain the expected GetDict method")

    generated_start = source.find("        // <auto-generated block-state-properties>")
    legacy_start = source.find("        private static readonly BlockStateDefinition[] stateDefinitions =")
    section_start = generated_start if generated_start >= 0 else legacy_start
    if section_start < 0:
        section_start = dictionary_index

    prefix = source[:section_start].rstrip()
    suffix = source[dictionary_index:]
    state_override = """        protected override BlockStateDefinition[] GetStateDefinitions()
        {
            return stateDefinitions;
        }
"""
    suffix = suffix.replace("\n" + state_override, "", 1)

    class_end = suffix.rfind("\n    }\n}")
    if class_end < 0:
        raise ValueError(f"{output_path} does not contain the expected class terminator")
    suffix = suffix[:class_end].rstrip() + "\n\n" + state_override.rstrip() + suffix[class_end:]

    output_path.write_text(
        prefix + "\n\n" + "\n".join(state_lines) + "\n\n" + suffix,
        encoding="utf-8",
    )
    return True


def main():
    if len(sys.argv) != 3:
        print(__doc__)
        sys.exit(1)

    blocks_json = Path(sys.argv[1])
    class_suffix = sys.argv[2]

    if not blocks_json.exists():
        print(f"Error: {blocks_json} not found")
        sys.exit(1)

    with blocks_json.open(encoding="utf-8") as f:
        data = json.load(f)

    # Build block ranges and compact state-property definitions, sorted by min state.
    block_ranges = []
    for block_key, block_info in data.items():
        cs_name = mc_name_to_csharp(block_key)
        states = block_info.get("states", [])
        state_ids = [s["id"] for s in states]
        if state_ids:
            properties = get_state_property_definitions(block_key, states, block_info.get("properties", {}))
            block_ranges.append((min(state_ids), max(state_ids), cs_name, properties))

    block_ranges.sort(key=lambda x: x[0])
    print(f"Loaded {len(block_ranges)} blocks from {blocks_json}")

    max_state = max(r[1] for r in block_ranges)
    print(f"State ID range: 0 - {max_state}")

    known_materials = load_known_materials()
    missing = [cs for _, _, cs, _ in block_ranges if known_materials and cs not in known_materials]
    if missing:
        print(f"\nWARNING: {len(missing)} blocks not found in Material.cs enum:")
        for cs_name in missing:
            print(f"  {cs_name}")
        print("\nYou need to add these to Material.cs before the palette will compile.")
        print("Insert them in alphabetical order within the enum.")

    class_name = f"Palette{class_suffix}"
    output_path = OUTPUT_DIR / OUTPUT_FILE_ALIASES.get(class_suffix, f"{class_name}.cs")
    state_lines, property_definition_count = render_state_definitions(block_ranges)

    if update_existing_palette(output_path, state_lines):
        print(
            f"Updated {output_path} with {property_definition_count} property definitions "
            "while preserving material mappings"
        )
        return

    lines = [
        "using System.Collections.Generic;",
        "",
        "namespace MinecraftClient.Mapping.BlockPalettes",
        "{",
        f"    public class {class_name} : BlockPalette",
        "    {",
        "        private static readonly Dictionary<int, Material> materials = new();",
        "",
        f"        static {class_name}()",
        "        {",
    ]

    for min_s, max_s, cs_name, _ in block_ranges:
        lines.append(f"            for (int i = {min_s}; i <= {max_s}; i++)")
        lines.append(f"                materials[i] = Material.{cs_name};")

    lines += [
        "        }",
        "",
        *state_lines,
        "",
    ]
    lines += [
        "        protected override Dictionary<int, Material> GetDict()",
        "        {",
        "            return materials;",
        "        }",
        "",
        "        protected override BlockStateDefinition[] GetStateDefinitions()",
        "        {",
        "            return stateDefinitions;",
        "        }",
        "    }",
        "}",
        "",
    ]

    output_path.write_text("\n".join(lines), encoding="utf-8")
    print(
        f"Generated {output_path} with {len(block_ranges)} blocks, "
        f"{max_state + 1} total states, and {property_definition_count} property definitions"
    )


if __name__ == "__main__":
    main()