#!/usr/bin/env python3
"""Write data-pack overrides so that structures place no ore.

Ore: removing ore features is not enough. Some structure templates have ore or
raw ore blocks built in (Cobblemon habitats and ruins), and some processor lists
turn blocks into ore (vanilla diamond fossils). For every template pool element
or configured feature (vanilla fossils, Cobblemon's `cobblemon:structure`) that
places such a template, and every processor list that outputs ore, this appends
one rule processor at the end of the processors that run for it, turning each
ore block into plain stone (see replacement()).

Vanilla mobs that structures place are kept out by the companion mod instead
(world/OnlyPokemon.kt), so structures keep all their pieces.

Data comes from the vanilla server jar and every mod jar in build/server/mods.
Re-run after a mod update and commit the result.

Usage: tools/clean_structures.py
"""
import copy
import gzip
import json
import re
import shutil
import sys
import zipfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import anvil  # noqa: E402
from mcserver import PACK, ROOT, SERVER  # noqa: E402

ORE = re.compile(r"(_ore|:raw_(iron|copper|gold)_block|:ancient_debris)$")
OUT = ROOT / "datapacks/core/data"
STRIP_ID = "cobblesurvival:strip_ores"
# Placeholder look-alikes for the raw ore blocks (decided in Phase 1, see docs/balance.md).
RAW = {"minecraft:raw_iron_block": "minecraft:tuff", "minecraft:raw_copper_block": "minecraft:terracotta",
       "minecraft:raw_gold_block": "minecraft:yellow_terracotta"}


def replacement(block):
    if block in RAW:
        return RAW[block]
    name = block.split(":", 1)[1]
    if name.startswith("deepslate_"):
        return "minecraft:deepslate"
    if name.startswith("nether_") or block == "minecraft:ancient_debris":
        return "minecraft:netherrack"
    return "minecraft:stone"


def sources():
    version = PACK["versions"]["minecraft"]
    jars = sorted(SERVER.glob(f"libraries/net/minecraft/server/{version}-*/server-{version}-*-extra.jar"))
    jars += sorted(p for p in (SERVER / "mods").glob("*.jar") if not p.name.startswith("cobblesurvival-"))
    if not jars:
        sys.exit("build/server is not set up; run tools/smoke_test.py once first")
    return [zipfile.ZipFile(j) for j in jars]


def resource_id(name, kind):
    """data/<ns>/<kind>/<path>.<ext> -> ns:path"""
    _, ns, rest = name.split("/", 2)
    return f"{ns}:{rest[len(kind) + 1:].rsplit('.', 1)[0]}"


def strip_processor(blocks):
    """A rule processor turning each ore block into stone, wrapped in `minecraft:capped`.

    Some lists use `minecraft:capped`, which changes blocks in a second, batch
    pass after every per-block processor, so a plain rule at the end can be
    undone. A capped processor with an unreachable cap runs its rule on every
    block in that batch pass, after the other processors in the list.
    """
    rules = [{"input_predicate": {"predicate_type": "minecraft:block_match", "block": b},
              "location_predicate": {"predicate_type": "minecraft:always_true"},
              "output_state": {"Name": replacement(b)}} for b in sorted(blocks)]
    return {"processor_type": "minecraft:capped", "limit": 1_000_000,
            "delegate": {"processor_type": "minecraft:rule", "rules": rules}}


def ore_outputs(node):
    """Ore blocks a processor (or anything nested in it, e.g. a capped delegate) can output."""
    if isinstance(node, dict):
        name = node.get("output_state", {}).get("Name", "") if isinstance(node.get("output_state"), dict) else ""
        found = {name} if ORE.search(name) else set()
        return found.union(*(ore_outputs(v) for v in node.values()))
    if isinstance(node, list):
        return set().union(*(ore_outputs(v) for v in node)) if node else set()
    return set()


def main():
    templates, processor_lists, pools, features = {}, {}, {}, {}
    for z in sources():
        for name in z.namelist():
            if not name.startswith("data/"):
                continue
            if name.endswith(".nbt") and "/structure/" in name:
                root = anvil.parse_nbt(gzip.decompress(z.read(name)))
                palette = root.get("palette") or [p for ps in root.get("palettes", []) for p in ps]
                ores = {p["Name"] for p in palette if ORE.search(p["Name"])}
                if ores:
                    templates[resource_id(name, "structure")] = ores
            elif name.endswith(".json") and "/worldgen/processor_list/" in name:
                processor_lists[resource_id(name, "worldgen/processor_list")] = json.loads(z.read(name))
            elif name.endswith(".json") and "/worldgen/template_pool/" in name:
                pools[resource_id(name, "worldgen/template_pool")] = json.loads(z.read(name))
            elif name.endswith(".json") and "/worldgen/configured_feature/" in name:
                features[resource_id(name, "worldgen/configured_feature")] = json.loads(z.read(name))

    ore_blocks = set().union(*templates.values()) if templates else set()
    lists_to_strip = set()
    for pid, plist in processor_lists.items():
        outputs = ore_outputs(plist)
        if outputs:
            ore_blocks |= outputs
            lists_to_strip.add(pid)
    for feature in features.values():
        config = feature.get("config")
        if not isinstance(config, dict):
            continue
        named = [v for value in config.values() for v in (value if isinstance(value, list) else [value])
                 if isinstance(v, str)]
        if any(t in templates for t in named):
            lists_to_strip.update(v for k, v in config.items() if "processors" in k and isinstance(v, str))

    pools_to_write = {}

    def visit(element, pool_id):
        if element.get("element_type") == "minecraft:list_pool_element":
            for inner in element["elements"]:
                visit(inner, pool_id)
            return
        if element.get("location") not in templates:
            return
        procs = element.get("processors", "minecraft:empty")
        if isinstance(procs, str) and procs != "minecraft:empty":
            lists_to_strip.add(procs)
        else:
            pools_to_write[pool_id] = True
            element["processors"] = STRIP_ID if procs == "minecraft:empty" else \
                {"processors": procs.get("processors", []) + [strip_processor(ore_blocks)]}

    edited_pools = {}
    for pid, pool in pools.items():
        copy_ = copy.deepcopy(pool)
        for entry in copy_.get("elements", []):
            visit(entry["element"], pid)
        if pid in pools_to_write:
            edited_pools[pid] = copy_

    strip = strip_processor(ore_blocks)
    written = []

    def write(kind, rid, data):
        ns, path = rid.split(":", 1)
        target = OUT / ns / kind / f"{path}.json"
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text(json.dumps(data, indent=2) + "\n")
        written.append(target)

    # Overrides always start from the original jar data, so clear the last run's output first.
    for kind in ("worldgen/processor_list", "worldgen/template_pool"):
        for ns_dir in OUT.iterdir():
            old = ns_dir / kind
            if old.is_dir() and ns_dir.name != "cobblesurvival":
                shutil.rmtree(old)
    write("worldgen/processor_list", STRIP_ID, {"processors": [strip]})
    for pid in sorted(lists_to_strip):
        if pid not in processor_lists:
            print(f"warning: processor list {pid} not found")
            continue
        procs = processor_lists[pid].get("processors", [])
        write("worldgen/processor_list", pid, {"processors": procs + [strip]})
    for pid, pool in sorted(edited_pools.items()):
        write("worldgen/template_pool", pid, pool)

    print(f"{len(templates)} templates contain ore; {len(ore_blocks)} ore blocks get replaced")
    print(f"wrote {len(written)} files ({len(lists_to_strip)} processor lists, {len(edited_pools)} template pools)")
    for block in sorted(ore_blocks):
        print(f"  {block} -> {replacement(block)}")


if __name__ == "__main__":
    main()