#!/usr/bin/env python3
"""Give every zone its own Pokémon: write one spawn pool per zone from Cobblemon's own spawn data.

Cobblemon's spawn files are switched off for the island (datapacks/core/pack.mcmeta filters
data/cobblemon/spawn_pool_world/). In their place this tool writes
datapacks/core/data/cobblesurvival/spawn_pool_world/<zone>.json:

  1. Every evolution line that Cobblemon lets spawn somewhere on the island gets one home
     zone. Lines with the fewest possible biomes are placed first. A line's score for a zone
     is how much of the zone is made of the line's biomes, times how much of the line's own
     range the zone covers (square root), times how well the line's strength
     (base-stat total of its strongest form, as a rank among all lines) matches the zone's
     level band, divided by how full the zone already is against its "lines" target. So strong
     lines (pseudo-legendaries) end up in high-level zones. Legendary, mythical, Ultra Beast
     and Paradox Pokémon are left out. Lines that can spawn in nearly every biome (Cobblemon's
     catch-all tags) are placed last: first into the cave zones, which take those whose first stage's
     primary type fits the cave (world/zones.json → caves → "types") up to their target, the ones closest to the
     cave's level first; the rest by strength into the land zones, weakest first, filling each zone
     up to its target.
  2. In its home zone a line keeps Cobblemon's spawn entries (bucket, weight, time, light,
     presets) for the biomes the zone has, plus "zone": "<id>". Levels are clamped into the
     zone's band, so the band picks the evolution stages: a stage whose levels miss the band
     doesn't spawn there. If no stage fits, the closest one spawns across the whole band. Every
     stage that reaches the band spawns up to its top (Adrian, 2026-10-31; Cobblemon's own ranges
     stop where the next stage takes over, which left e.g. Dusclops at Lv 45–46 in the Deep Dark).
  3. world/zones.json → "species" pins a line to a zone by hand: {"pikachu": "southwold"}. The long form
     {"ditto": {"zone": "deep_dark", "all_entries": true}} also keeps all of the line's entries when no stage
     fits the band (otherwise the closest stage brings one entry).
  4. In cave zones every entry spawns in the cave's biome, and sky-light minimums, "can see sky", time
     of day, excluded biomes and the "natural" ground preset are dropped (there's no sky or day underground).
  5. Strong flyers (STRONG_FLYER_*) only spawn in zones whose band starts at HIGH_ZONE_LEVEL or more.
  6. Each zone with a "guardian" gets datapacks/core/data/cobblesurvival/guardians/<zone>.json: the
     strongest Pokémon spawning in the zone (or "species"), uncatchable, half again as big, five levels
     above the band, with the zone's boss drops and world/zones.json → guardian_respawn_hours.

Also writes docs/zone-species.md, the list per zone for review.

Needs: build/zone-layout.json (tools/zone_layout.py), build/field-guide/fieldguide.json
(tools/fieldguide/build.py; its "hidden_spawns" say which island biomes each of Cobblemon's own,
filtered-out entries can use) and the
Cobblemon jar in build/server/mods (any test-server run).

Usage: tools/zone_spawns.py
"""
import collections
import copy
import json
import re
import sys
import zipfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "datapacks/core/data/cobblesurvival/spawn_pool_world"
GUARDIANS = ROOT / "datapacks/core/data/cobblesurvival/guardians"
GUARDIAN_SCALE = 1.5  # "Alpha": half again as big as a normal one
GUARDIAN_LEVELS_ABOVE_BAND = 5
REPORT = ROOT / "docs/zone-species.md"
EXCLUDED_LABELS = {"legendary", "mythical", "ultra_beast", "paradox", "restricted"}
MIN_CELLS = 5  # a zone "has" a biome when at least this many grid cells of it lie in the zone
# A line whose entries cover this share of the island's biomes lives "anywhere" (Cobblemon's catch-all
# tags): its biomes say nothing about a home, so it only goes to a land zone, by strength. Its cave and
# water entries still spawn underground and in the rivers of that zone.
ALL_ROUNDER = 0.8
# Roadmap Phase 3: strong flyers only spawn in high-level zones. A species you can ride through the air with a top
# flying speed or stamina this high never spawns in a zone whose band starts below HIGH_ZONE_LEVEL (its weaker
# pre-evolutions still do), and a line that starts as one is placed in such a zone only; if none fits its habitat,
# it doesn't spawn in the wild.
STRONG_FLYER_SPEED = 60
STRONG_FLYER_STAMINA = 85
HIGH_ZONE_LEVEL = 30


def cobblemon_jar():
    jars = sorted((ROOT / "build/server/mods").glob("Cobblemon-*.jar"))
    if not jars:
        sys.exit("no Cobblemon jar in build/server/mods: run tools/build.py and tools/selftest.py (or the quick check) first")
    return jars[-1]


def load_entries(jar):
    """{(pool, id): entry} for every spawn entry of Cobblemon's world spawn pools."""
    entries = {}
    with zipfile.ZipFile(jar) as z:
        for name in z.namelist():
            m = re.fullmatch(r"data/cobblemon/spawn_pool_world/(.+)\.json", name)
            if not m:
                continue
            data = json.loads(z.read(name))
            if not data.get("enabled", True) or data.get("neededInstalledMods"):
                continue
            pool = f"cobblemon:{m[1]}"
            for entry in data.get("spawns", []):
                # Same key as the field guide, which names an entry without an id after its file.
                entries[(pool, entry.get("id") or pool)] = entry
    return entries


def zone_biomes(plan, layout, preview):
    """{zone: {biome: share of the zone's cells}}, ignoring biomes with fewer than MIN_CELLS cells."""
    names = preview["biomes"]
    counts = collections.defaultdict(collections.Counter)
    for zone_index, biome in zip(layout["grid"], preview["surface"]):
        if zone_index >= 0:
            counts[layout["zones"][zone_index]][names[biome]] += 1
    shares = {}
    for zid, c in counts.items():
        kept = {b: n for b, n in c.items() if n >= MIN_CELLS}
        total = sum(kept.values())
        shares[zid] = {b: n / total for b, n in kept.items()}
    for cave in plan["caves"]:
        shares[cave["id"]] = {cave["biome"]: 1.0}
    return shares


def species_data(jar):
    """({species id: base-stat total}, {species ids that are strong flyers}) from Cobblemon's species files."""
    totals, flyers = {}, set()
    with zipfile.ZipFile(jar) as z:
        for name in z.namelist():
            m = re.fullmatch(r"data/cobblemon/species/[^/]+/([^/]+)\.json", name)
            if not m:
                continue
            data = json.loads(z.read(name))
            totals[f"cobblemon:{m[1]}"] = sum(data.get("baseStats", {}).values())
            air = ((data.get("riding") or {}).get("behaviours") or {}).get("AIR") or {}
            stats = air.get("stats", {})
            def top(stat):
                return int(str(stats.get(stat, "0")).split("-")[-1])
            if air and (top("SPEED") >= STRONG_FLYER_SPEED or top("STAMINA") >= STRONG_FLYER_STAMINA):
                flyers.add(f"cobblemon:{m[1]}")
    return totals, flyers


def load_before(zone, room, stretch, target):
    """How many lines a land zone may hold once all-rounders are added."""
    return target[zone] - room[zone] + room[zone] * stretch


def line_root(species_id, species):
    seen = set()
    while species_id in species and species[species_id].get("pre") and species_id not in seen:
        seen.add(species_id)
        species_id = species[species_id]["pre"]
    return species_id


def main():
    plan = json.loads((ROOT / "world/zones.json").read_text())
    layout = json.loads((ROOT / "build/zone-layout.json").read_text())
    recipe = json.loads((ROOT / "world/world.json").read_text())
    preview = json.loads((ROOT / f"build/server/seed-search/cobblesurvival/biome_preview/{recipe['seed']}.json").read_text())
    guide = json.loads((ROOT / "build/field-guide/fieldguide.json").read_text())
    species = guide["species"]
    jar = cobblemon_jar()
    entries = load_entries(jar)
    bst, strong_flyers = species_data(jar)

    zones = {z["id"]: z for z in plan["zones"]} | {c["id"]: c for c in plan["caves"]}
    shares = zone_biomes(plan, layout, preview)
    biomes_of = {zid: set(s) for zid, s in shares.items()}
    island_biomes = set().union(*biomes_of.values())
    land_zones = {z["id"] for z in plan["zones"] if "water" not in z["rule"]}
    cave_zones = {c["id"] for c in plan["caves"]}
    mids = {zid: sum(z["levels"]) / 2 for zid, z in zones.items()}
    low, high = min(mids.values()), max(mids.values())
    tier = {zid: (m - low) / (high - low) for zid, m in mids.items()}

    # Island biomes of every Cobblemon entry, from the field guide (it resolves biome tags). The guide lists
    # pools that datapacks/core's filter hides, Cobblemon's own, under "hidden_spawns"; older guides had
    # them in "spawns".
    cobblemon_spawns = guide.get("hidden_spawns") or guide["spawns"]
    lines = collections.defaultdict(list)  # line root -> [(entry, island biomes, species id)]
    for s in cobblemon_spawns:
        if s["type"] != "pokemon" or s["status"] != "available" or not s["biomes"]:
            continue
        entry = entries.get((s["pool"], s["id"]))
        if entry is None:
            continue
        member = s["members"][0]["species"]
        if EXCLUDED_LABELS & set(species.get(member, {}).get("labels", [])):
            continue
        lines[line_root(member, species)].append((entry, set(s["biomes"]), member, s["level"]))

    pins = {f"cobblemon:{k}" if ":" not in k else k: v for k, v in plan.get("species", {}).items()}
    pinned = {root: pin["zone"] if isinstance(pin, dict) else pin for root, pin in pins.items()}
    all_entries = {root for root, pin in pins.items() if isinstance(pin, dict) and pin.get("all_entries")}
    flying_lines = {root for root in lines if root in strong_flyers}

    def allowed(root, zid):
        return root not in flying_lines or zones[zid]["levels"][0] >= HIGH_ZONE_LEVEL
    wanted = {zid: z.get("lines", 20) for zid, z in zones.items()}
    scale = len(lines) / sum(wanted.values())
    target = {zid: n * scale for zid, n in wanted.items()}
    # Strength of a line: the base-stat total of its strongest member, as a rank from 0 (weakest) to 1.
    strength = {root: max(bst.get(m, 0) for _, _, m, _ in members) for root, members in lines.items()}
    ranked = sorted(strength, key=strength.get)
    rank = {root: i / max(len(ranked) - 1, 1) for i, root in enumerate(ranked)}
    load = collections.Counter()
    home = {}
    biomes_of_line = {root: set().union(*(b for _, b, _, _ in members)) for root, members in lines.items()}
    all_rounders = {root for root, b in biomes_of_line.items()
                    if root not in pinned and len(b) >= ALL_ROUNDER * len(island_biomes)}
    for root, zid in pinned.items():
        if root in lines:
            home[root] = zid
            load[zid] += 1

    # 1. Lines with a habitat, the most specific first.
    for root in sorted(lines, key=lambda r: (len(biomes_of_line[r]), r)):
        if root in home or root in all_rounders:
            continue
        line_biomes = biomes_of_line[root]

        def score(zid):
            habitat = sum(share for b, share in shares[zid].items() if b in line_biomes)
            # How much of the line's own range this zone covers.
            specific = len(line_biomes & biomes_of[zid]) / len(line_biomes)
            fit = 1 - abs(rank[root] - tier[zid])
            return habitat * specific ** 0.5 * fit ** 2 / (1 + load[zid] / target[zid])

        candidates = [zid for zid in zones if line_biomes & biomes_of[zid] and allowed(root, zid)]
        if candidates:
            home[root] = max(candidates, key=score)
            load[home[root]] += 1

    # 2. Cave zones take the all-rounders of the types that fit them, up to their targets, the ones closest
    #    to the cave's level first (Adrian, 2026-10-30: the deep caves' tables were too small).
    # A line's type here is its first stage's primary type (a line of many types, like Eevee's, shouldn't count as all of them).
    primary_type = {root: (species.get(root, {}).get("types") or [None])[0] for root in lines}
    for zid in sorted(cave_zones, key=lambda z: tier[z]):
        fits = set(zones[zid].get("types", []))
        while load[zid] < target[zid]:
            pool = [r for r in all_rounders if r not in home and primary_type[r] in fits and allowed(r, zid)]
            if not pool:
                break
            root = min(pool, key=lambda r: (abs(rank[r] - tier[zid]), r))
            home[root] = zid
            load[zid] += 1

    # 3. The other all-rounders fill the land zones up to their targets, weakest lines into the lowest zones.
    by_tier = sorted(land_zones, key=lambda z: tier[z])
    room = {z: max(0.0, target[z] - load[z]) for z in by_tier}
    total_room = sum(room.values()) or 1
    stretch = max(1.0, len(all_rounders) / total_room)
    zone_iter = iter(by_tier)
    current = next(zone_iter)
    for root in sorted((r for r in all_rounders if r not in home), key=lambda r: rank[r]):
        while load[current] >= load_before(current, room, stretch, target) and current != by_tier[-1]:
            current = next(zone_iter)
        # A strong flyer skips ahead to the first high zone.
        home[root] = current if allowed(root, current) else next((z for z in by_tier if allowed(root, z)), None)
        if home[root] is None:
            del home[root]
            continue
        load[home[root]] += 1

    pools = collections.defaultdict(list)
    report = collections.defaultdict(list)
    cave_biome = {c["id"]: c["biome"] for c in plan["caves"]}
    for root, zid in home.items():
        lo, hi = zones[zid]["levels"]
        high = zones[zid]["levels"][0] >= HIGH_ZONE_LEVEL
        # Underground every entry gets the cave's biome below, so any of the line's entries will do.
        usable = [(e, b, m, lv) for e, b, m, lv in lines[root]
                  if (zid in cave_zones or b & biomes_of[zid]) and (high or m not in strong_flyers)]
        if not usable:
            continue
        # Every stage that reaches the band spawns up to its top; Cobblemon's own ranges stop where the next stage
        # takes over. A caught one past its evolution level evolves on its next level-up.
        fitting = [(e, m, max(lv[0], lo), hi) for e, b, m, lv in usable if lv[0] <= hi and lv[1] >= lo]
        if not fitting:
            # No stage reaches the band: the closest stage spawns across all of it.
            e, b, m, lv = min(usable, key=lambda u: min(abs(u[3][0] - hi), abs(u[3][1] - lo)))
            fitting = [(e, m, lo, hi)]
            if root in all_entries:
                fitting = [(ue, m, lo, hi) for ue, _, um, ulv in usable if um == m and ulv == lv]
        seen = set()
        for entry, member, a, b in fitting:
            spawn = copy.deepcopy(entry)
            spawn["id"] = f"{zid}-{entry.get('id') or member.split(':')[1]}-{len(pools[zid])}"
            spawn["level"] = f"{a}-{b}"
            condition = spawn.setdefault("condition", {})
            condition["zone"] = zid
            if zid in cave_zones:
                # Many entries come from Cobblemon's surface spawns (e.g. "spooky" biomes include the deep
                # dark); underground they could never see the sky.
                condition.pop("minSkyLight", None)
                if condition.get("canSeeSky") is True:
                    del condition["canSeeSky"]
                # "natural" only allows natural ground; cave floors (basalt deltas: magma, smooth basalt) often aren't.
                if "natural" in spawn.get("presets", []):
                    spawn["presets"] = [x for x in spawn["presets"] if x != "natural"]
                # The zone condition places it; the cave's biome stands in for Cobblemon's surface biomes, and
                # there's no day or night underground.
                condition["biomes"] = [cave_biome[zid]]
                condition.pop("timeRange", None)
                spawn.get("anticondition", {}).pop("biomes", None)
                if spawn.get("anticondition") == {}:
                    del spawn["anticondition"]
                key = (spawn["pokemon"], spawn.get("spawnablePositionType"), spawn.get("bucket"), spawn["level"])
                if key in seen:
                    continue
                seen.add(key)
            pools[zid].append(spawn)
        names = sorted({species[m]["name"] for _, m, _, _ in fitting})
        report[zid].append((species.get(root, {}).get("name", root), names))

    low_flyers = [(zid, s["pokemon"]) for zid, spawns in pools.items() for s in spawns
                  if zones[zid]["levels"][0] < HIGH_ZONE_LEVEL and f"cobblemon:{s['pokemon'].split()[0].lower()}" in strong_flyers]
    if low_flyers:
        sys.exit(f"strong flyers would spawn below Lv {HIGH_ZONE_LEVEL}: {low_flyers[:5]}")

    OUT.mkdir(parents=True, exist_ok=True)
    for old in OUT.glob("*.json"):
        old.unlink()
    for zid, spawns in pools.items():
        (OUT / f"{zid}.json").write_text(json.dumps({
            "_comment": "Generated by tools/zone_spawns.py from Cobblemon's spawn data; do not edit by hand.",
            "enabled": True, "neededInstalledMods": [], "neededUninstalledMods": [],
            "spawns": sorted(spawns, key=lambda s: s["id"]),
        }, indent=1) + "\n")

    GUARDIANS.mkdir(parents=True, exist_ok=True)
    for old in GUARDIANS.glob("*.json"):
        old.unlink()
    guardians = {}
    for zid, zone in zones.items():
        config = zone.get("guardian")
        if not config:
            continue
        # A guardian stands its ground: on land a walking spawn, at sea a swimming one, never a fishing catch.
        kinds = {"submerged", "surface"} if "water" in zone.get("rule", {}) else {"grounded"}
        spawning = {s["pokemon"].split()[0].lower() for s in pools.get(zid, []) if s.get("spawnablePositionType") in kinds}
        species_name = config.get("species") or max(spawning, key=lambda n: bst.get(f"cobblemon:{n}", 0))
        level = zone["levels"][1] + GUARDIAN_LEVELS_ABOVE_BAND
        guardians[zid] = species_name
        (GUARDIANS / f"{zid}.json").write_text(json.dumps({
            "_comment": "Generated by tools/zone_spawns.py from world/zones.json; do not edit by hand.",
            "zone": zid,
            "pokemon": f"{species_name} level={level} uncatchable=yes",
            "scale": GUARDIAN_SCALE,
            "drops": [{"item": item, "count": count} for item, count in config["drops"]],
            "respawn_hours": plan["guardian_respawn_hours"],
        }, indent=1) + "\n")

    lines_out = ["# Pokémon per zone", "",
                 "Generated by `tools/zone_spawns.py` from Cobblemon's spawn data and `world/zones.json`; "
                 "do not edit by hand. Each evolution line spawns in one zone; the stages listed are the ones "
                 "whose levels reach the zone's band. Pin a line to another zone in `world/zones.json` → "
                 "`species`.", ""]
    for zid, zone in zones.items():
        rows = sorted(report.get(zid, []))
        lines_out += [f"## {zone['name']} (Lv {zone['levels'][0]}–{zone['levels'][1]}, {len(rows)} lines)", ""]
        lines_out += [", ".join(f"{root}" + (f" ({', '.join(n)})" if n != [root] else "") for root, n in rows) or "none", ""]
    lines_out += ["## Alpha guardians", "",
                  f"One per zone at its sample point, uncatchable, five levels above the band; it returns "
                  f"{plan['guardian_respawn_hours']} hours after it is beaten. Its drops help with the next zone.", "",
                  "| Zone | Guardian | Level | Drops |", "| --- | --- | --- | --- |"]
    for zid, species_name in guardians.items():
        zone = zones[zid]
        drop_text = ", ".join(f"{count}× {item.split(':')[1]}" for item, count in zone["guardian"]["drops"])
        lines_out.append(f"| {zone['name']} | {species[f'cobblemon:{species_name}']['name'] if f'cobblemon:{species_name}' in species else species_name} "
                         f"| {zone['levels'][1] + GUARDIAN_LEVELS_ABOVE_BAND} | {drop_text} |")
    REPORT.write_text("\n".join(lines_out) + "\n")

    unplaced = [r for r in lines if r not in home]
    grounded = sorted(species.get(r, {}).get("name", r) for r in unplaced if r in flying_lines)
    if grounded:
        print(f"strong flyers kept out of the wild (no high zone fits their habitat): {', '.join(grounded)}")
    for zid, zone in zones.items():
        print(f"  {zone['name']:18s} {len(report.get(zid, [])):3d} lines  {len(pools.get(zid, [])):4d} spawn entries")
    print(f"{len(home)} lines placed, {len(unplaced)} with no matching zone; wrote {OUT.relative_to(ROOT)} and {REPORT.relative_to(ROOT)}")


if __name__ == "__main__":
    main()