evolution.py

"""Evolutions and pre-evolutions of existing dex species.

The dex owns the chain rules (at most 3 stages, no cycles, legendaries
neither evolve nor are evolved from - tools/dexlib.py there re-checks every
submission at PR time), but a rejection that only arrives after the artwork
is paid for is a bad experience. This module answers the two questions the
Create tab must answer before spending anything: is this anchor allowed, and
what family context does the new stage inherit?

It reads the dex snapshot's species.json/spec.json directly, not the derived
entries list: those are the files the dex validates, so the guard and the
gatekeeper read the same fields.
"""

import json
import os

from fakemon_forge import species as species_mod

MAX_CHAIN_LENGTH = 3  # mirrors tools/dexlib.py in the dex repo

METHODS = {
    "Level up": "EVO_LEVEL",
    "Trade": "EVO_TRADE",
    "Friendship": "EVO_FRIENDSHIP",
}


def _load(path):
    try:
        with open(path) as f:
            return json.load(f)
    except (OSError, ValueError):
        return {}


def chains(dex_root):
    """Per-species evolution context, keyed by directory slug.

    Each entry carries what the guard needs (links, legendary flags) and
    what the new stage inherits (family, family seed, stage position, the
    anchor's committed BST, its raw artwork for visual continuity).
    """
    base = os.path.join(dex_root, "dex")
    out = {}
    if not os.path.isdir(base):
        return out
    for slug in sorted(os.listdir(base)):
        sdir = os.path.join(base, slug)
        if not os.path.isdir(sdir):
            continue
        spec = _load(os.path.join(sdir, "spec.json"))
        species = _load(os.path.join(sdir, "species.json"))
        if not spec and not species:
            continue
        targets = [l.get("target") for l in (species.get("evolvesTo") or [])
                   if isinstance(l, dict) and l.get("target")]
        stats = [species.get(k) or 0 for k in
                 ("baseHP", "baseAttack", "baseDefense",
                  "baseSpAttack", "baseSpDefense", "baseSpeed")]
        out[slug] = {
            "name": spec.get("name") or species.get("speciesName") or slug,
            "family": spec.get("family") or slug,
            "family_seed": spec.get("family_seed", spec.get("seed", 0)),
            "stage": int(spec.get("stage") or 1),
            "of_stages": int(spec.get("of_stages") or 1),
            "from": species.get("evolvesFrom"),
            "to": targets,
            "legendary": bool(species.get("isSubLegendary") or
                              species.get("isRestrictedLegendary")),
            "bst": sum(s for s in stats if isinstance(s, int)),
            "raw": os.path.join(sdir, "raw_front.png"),
        }
    return out


def _back_len(ch, slug):
    """Stages before `slug` on its chain. Cycle-safe: corrupt data ends the walk."""
    seen, n, total = set(), slug, 0
    while True:
        prev = ch.get(n, {}).get("from")
        if not prev or prev in seen or prev not in ch:
            return total
        seen.add(prev)
        n, total = prev, total + 1


def _fwd_len(ch, slug):
    """Longest forward path from `slug`. Cycle-safe."""
    best, stack = 0, [(slug, 0, {slug})]
    while stack:
        n, d, seen = stack.pop()
        best = max(best, d)
        for t in ch.get(n, {}).get("to", []):
            if t not in seen and t in ch:
                stack.append((t, d + 1, seen | {t}))
    return best


def check_anchor(ch, anchor, direction):
    """None when the anchor may take this addition, else the reason it may not.

    direction is "evolution" (the new stage evolves from the anchor) or "pre"
    (the anchor evolves from the new stage). The dex re-checks all of this at
    PR time; this is the friendly front stop before paid generation.
    """
    info = ch.get(anchor)
    if info is None:
        return f"Unknown anchor {anchor!r} - reload the dex and try again."
    name = info["name"]
    if info["legendary"]:
        return (f"{name} is legendary - legendaries cannot evolve or have "
                "pre-evolutions.")
    if direction == "pre":
        if info["from"]:
            return f"{name} already has a pre-evolution."
        if _fwd_len(ch, anchor) + 2 > MAX_CHAIN_LENGTH:
            return (f"A pre-evolution for {name} would make a chain longer "
                    f"than {MAX_CHAIN_LENGTH}.")
    elif _back_len(ch, anchor) + 2 > MAX_CHAIN_LENGTH:
        return (f"{name} is already at the end of a {MAX_CHAIN_LENGTH}-stage "
                "chain.")
    return None


def new_stage_context(ch, anchor, direction):
    """Family context and BST for the new stage. check_anchor() must pass first.

    The new stage joins the anchor's family: same family seed (so the battle
    style and stat-shape draws match the line), a stage position past the
    anchor - or before it, for a pre-evolution - and a BST drawn from the
    family's ladder as the rung it would occupy, clamped so an evolution is
    always an upgrade over the anchor's committed stats and a pre-evolution
    always a downgrade. The anchor's BST is real data; the re-drawn ladder is
    hypothetical, so the clamp wins ties.
    """
    info = ch[anchor]
    seed_spec = {"family_seed": info["family_seed"]}
    if direction == "pre":
        stage = 1
        of_stages = _fwd_len(ch, anchor) + 2
        rung = species_mod.family_bst_ladder(seed_spec, of_stages)[0]
        total = max(species_mod.BST_STAGE_FLOOR,
                    min(rung, info["bst"] - species_mod.BST_STAGE_GAP))
    else:
        stage = info["stage"] + 1
        of_stages = _back_len(ch, anchor) + 2
        rung = species_mod.family_bst_ladder(seed_spec, stage)[-1]
        total = max(rung, info["bst"] + species_mod.BST_STAGE_GAP)
    # The legendary band is a ROM-level invariant; it outranks the upgrade gap
    # on the rare anchor whose committed BST sits right below it.
    total = min(total, species_mod.BST_LEGENDARY_MIN - 1)
    return {"family": info["family"], "family_seed": info["family_seed"],
            "stage": stage, "of_stages": of_stages, "total": total}