#!/usr/bin/env python3
"""Write the data behind the pack's own JEI pages (roadmap Phase 8; Adrian, 2026-11-02: every item that can drop gets
a JEI page naming its sources; 2026-11-03: alongside the Cobblemon Info add-on, which shows spawns and drops but not
the island's zones).

JEI runs on the client, which on a dedicated server doesn't have the pack's server data, so the server sends each player
this file (dev.cobblesurvival.jei.DropSources): datapacks/core/data/cobblesurvival/drop_sources/pages.json. It holds:

  pokemon        every species of a line that lives on the island: its drops (as the Drop Synthesizer rolls them)
                 and the zones it spawns in, or the zones where its line lives
                 (legendaries: where a Sigil summons them, world/summons.json; that line's evolutions too)
  summons        each Sigil: species, name, zone, level, so JEI can list it as an item of its own
  spawn_drops    zone spawns that drop something else than their species (e.g. Beheeyem in the End)
  guardians      each zone guardian: species, level, drops, return time
  end_boss       the End boss and what beating it gives
  altar          the Sacrifice Altar's level tiers and type specials; rebirth: the Rebirth Chamber's extra
  ore_veins      per zone, what its regenerating veins give and how fast they regrow
  pasture_finds  per Cobblebase skill and zone, what a pasture's Pokémon can find (the pack's zone-bound tables)

The self-test drop_sources_match_the_server compares it with the live server, so run this after changing species
drops, spawns, guardians, the End boss, the altar, ore veins or Cobblebase's tables (tools/build.py and
tools/server_update.py run it).

Usage: tools/jei_sources.py
"""
import json
import re
import zipfile
from collections import defaultdict
from pathlib import Path

import lock

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "datapacks/core/data"
OUT = DATA / "cobblesurvival/drop_sources/pages.json"
SKILLS = {"mining": "mining", "finder_ore": "ore_finder", "finder_evo": "evolution_item_finder"}


def amount_range(drops):
    a, _, b = str(drops.get("amount", 1)).partition("-")
    return range(int(a), int(b or a) + 1)


def yield_chance(entry):
    """The chance a picked entry gives at least one item: a count of 0 gives nothing."""
    low, _, high = str(entry.get("quantityRange", entry.get("quantity", 1))).partition("-")
    low, high = int(low), int(high or low)
    return 1.0 if low > 0 else high / (high + 1)


def selection_odds(entries, chosen):
    """{item: the chance one drop roll for [chosen] picks gives at least one}, as DropTable.getDrops works: in table
    order, the first entry that passes its percentage is picked; a round with none uses up a pick. A table can list
    one item twice (Snorlax's apples), so this follows items, not entries."""
    percent = [float(e.get("percentage", 100)) / 100 for e in entries]
    quantity = [int(e.get("quantity", 1)) for e in entries]
    limit = [int(e.get("maxSelectableTimes", 1)) for e in entries]
    start = (tuple(i for i in range(len(entries)) if quantity[i] <= chosen), 0, tuple(0 for _ in entries), frozenset())
    states, odds = {start: 1.0}, defaultdict(float)
    while states:
        following = defaultdict(float)
        def advance(possible, count, times, have, p):
            if count < chosen and possible:
                following[(possible, count, times, have)] += p
        for (possible, count, times, have), p in states.items():
            miss = 1.0
            for i in possible:
                hit = miss * percent[i]
                miss *= 1 - percent[i]
                if hit <= 0:
                    continue
                now = count + quantity[i]
                picked = times[:i] + (times[i] + 1,) + times[i + 1:]
                left = tuple(j for j in possible if not (j == i and picked[j] >= limit[j]) and quantity[j] <= chosen - now)
                item, gives = entries[i].get("item"), yield_chance(entries[i])
                if item is None or item in have:
                    advance(left, now, picked, have, p * hit)
                    continue
                odds[item] += p * hit * gives
                advance(left, now, picked, have | {item}, p * hit * gives)
                if gives < 1:
                    advance(left, now, picked, have, p * hit * (1 - gives))
            if miss > 0:
                advance(possible, count + 1, times, have, p * miss)
        states = following
    return odds


def drop_entries(drops):
    """Cobblemon drop table → [{item, chance (percent that a roll gives at least one, absent = always), count ("a-b")}]."""
    drops = drops or {}
    entries = drops.get("entries", [])
    picks = amount_range(drops)
    odds = defaultdict(float)
    if entries:
        for chosen in picks:
            for item, chance in selection_odds(entries, chosen).items():
                odds[item] += chance / len(picks)
    counts = {}
    for e in entries:
        if "item" in e:
            low, _, high = str(e.get("quantityRange", e.get("quantity", 1))).partition("-")
            low, high = max(int(low), 1), int(high or low)
            a, b = counts.get(e["item"], (low, high))
            counts[e["item"]] = (min(a, low), max(b, high))
    out = []
    for item, (low, high) in counts.items():
        entry = {"item": item, "count": str(low) if low == high else f"{low}-{high}"}
        if odds[item] < 0.9995:
            entry["chance"] = round(odds[item] * 100, 1)
        out.append(entry)
    return out


def species_data(jar):
    species = {}
    with zipfile.ZipFile(jar) as z:
        for name in z.namelist():
            m = re.fullmatch(r"data/cobblemon/species/[^/]+/([^/]+)\.json", name)
            if not m:
                continue
            d = json.loads(z.read(name))
            pre = d.get("preEvolution")
            species[f"cobblemon:{m[1]}"] = {
                "name": d["name"], "implemented": d.get("implemented", False), "drops": d.get("drops"),
                "pre": f"cobblemon:{pre.split()[0].lower()}" if pre else None,
            }
    # The pack's species additions replace a species' drops (datapacks/core/data/cobblesurvival/species_additions).
    for path in sorted((DATA / "cobblesurvival/species_additions").glob("*.json")):
        d = json.loads(path.read_text())
        if "drops" in d and d["target"] in species:
            species[d["target"]]["drops"] = d["drops"]
    return species


def root_of(sid, species):
    seen = set()
    while species.get(sid, {}).get("pre") and sid not in seen:
        seen.add(sid)
        sid = species[sid]["pre"]
    return sid


def main():
    species = species_data(lock.mod_jar("cobblemon"))
    plan = json.loads((ROOT / "world/zones.json").read_text())
    zones = {z["id"]: z for z in plan["zones"] + plan["caves"] + plan.get("dimensions", [])}

    # Where each species spawns, and spawns with drops of their own.
    where = defaultdict(dict)
    spawn_drops = []
    for zid in zones:
        pool = DATA / f"cobblesurvival/spawn_pool_world/{zid}.json"
        if not pool.exists():
            continue
        for spawn in json.loads(pool.read_text())["spawns"]:
            sid = f"cobblemon:{spawn['pokemon'].split()[0].lower()}"
            a, _, b = str(spawn["level"]).partition("-")
            low, high = where[sid].get(zid, (999, 0))
            where[sid][zid] = (min(low, int(a)), max(high, int(b or a)))
            if "drops" in spawn:
                entry = {"species": sid, "zone": zid, "drops": drop_entries(spawn["drops"])}
                if entry not in spawn_drops:
                    spawn_drops.append(entry)

    line_zones = defaultdict(dict)
    for sid, zs in where.items():
        for zid, levels in zs.items():
            line_zones[root_of(sid, species)].setdefault(zid, sid)

    table = json.loads((ROOT / "world/summons.json").read_text())
    level_of = lambda v: table["tiers"][str(v["tier"])]["level"]
    summoned = {f"cobblemon:{k}": {"zone": v["zone"], "level": level_of(v)} for k, v in table["summons"].items() if "form" not in v}
    def sigil(key, v):
        base = f"cobblemon:{v.get('species', key)}"
        name = f"{v['form'].title()} {species[base]['name']}" if "form" in v else species[base]["name"]
        return {"species": f"cobblemon:{key}", "pokemon": base, "name": name, "zone": v["zone"], "level": level_of(v)}
    summons = [sigil(k, v) for k, v in table["summons"].items()]

    pokemon = []
    for sid, s in sorted(species.items()):
        drops = drop_entries(s["drops"])
        root = root_of(sid, species)
        if not drops or not s["implemented"] or not (where.get(sid) or line_zones.get(root) or root in summoned):
            continue
        entry = {"species": sid, "drops": drops,
                 "zones": [{"zone": z, "levels": f"{lo}-{hi}"} for z, (lo, hi) in sorted(where.get(sid, {}).items())]}
        if root in summoned:
            entry["summon"] = summoned[root]
        elif not entry["zones"]:
            entry["line"] = [{"zone": z, "species": member} for z, member in sorted(line_zones[root].items())]
        pokemon.append(entry)

    guardians = []
    for path in sorted((DATA / "cobblesurvival/guardians").glob("*.json")):
        g = json.loads(path.read_text())
        name, *props = g["pokemon"].split()
        guardians.append({"zone": g["zone"], "species": f"cobblemon:{name}",
                          "level": int(next(p.split("=")[1] for p in props if p.startswith("level="))),
                          "drops": g["drops"], "respawn_hours": g["respawn_hours"]})

    end = json.loads((DATA / "cobblesurvival/bosses/end.json").read_text())["end"]
    end_boss = {"species": f"cobblemon:{end['species']}", "level": end["level"], "xp_first_win": end["xp_first_win"],
                "xp": end["xp"], "respawn_hours": end["respawn_hours"],
                "first_win": [{"item": "minecraft:dragon_egg", "count": 1}]}

    machines = DATA / "cobblesurvival/machines"
    altar = json.loads((machines / "altar_core.json").read_text())["altar"]
    rebirth = json.loads((machines / "rebirth_core.json").read_text())["rebirth"]

    ore_veins = defaultdict(list)
    for ore, spec in plan["ores"].items():
        ore_veins[spec["zone"]].append({"item": spec["item"], "regen_minutes": spec["regen_minutes"]})

    finds = defaultdict(lambda: defaultdict(set))
    for path in sorted((DATA / "cobblebase/loot_table").glob("*.json")):
        prefix, _, tier = path.stem.rpartition("_")
        skill = SKILLS.get(prefix)
        if not skill:
            continue
        for pool in json.loads(path.read_text()).get("pools", []):
            for e in pool.get("entries", []):
                if e.get("type") != "minecraft:item":
                    continue
                zone_conditions = [c for c in e.get("conditions", []) if c.get("condition") == "cobblesurvival:in_zone"]
                for zid in (zone_conditions[0]["zones"] if zone_conditions else ["*"]):
                    finds[(skill, zid)][e["name"]].add(tier)
    pasture_finds = [{"skill": skill, "zone": None if zid == "*" else zid,
                      "items": [{"item": item, "tiers": sorted(tiers)} for item, tiers in sorted(items.items())]}
                     for (skill, zid), items in sorted(finds.items(), key=lambda kv: (kv[0][0], kv[0][1] != "*", kv[0][1]))]

    data = {
        "_comment": "Generated by tools/jei_sources.py from the pack's data; do not edit by hand. The server sends it to "
                    "players for the companion mod's JEI pages (jei/DropSources.kt); checked against the live server by the "
                    "self-test drop_sources_match_the_server.",
        "zones": {zid: {"name": z["name"], "levels": z["levels"]} for zid, z in zones.items()},
        "pokemon": pokemon, "summons": summons, "spawn_drops": spawn_drops, "guardians": guardians, "end_boss": end_boss,
        "altar": {"tiers": altar["tiers"], "specials": altar.get("specials", [])}, "rebirth": rebirth.get("extra", []),
        "ore_veins": [{"zone": z, "items": items} for z, items in sorted(ore_veins.items())],
        "pasture_finds": pasture_finds,
    }
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(data, indent=1, ensure_ascii=False) + "\n")
    print(f"{len(pokemon)} Pokémon, {len(spawn_drops)} spawns with own drops, {len(guardians)} guardians, "
          f"{len(ore_veins)} zones with veins, {len(pasture_finds)} pasture find lists; wrote {OUT.relative_to(ROOT)}")


if __name__ == "__main__":
    main()