"""
Part library generated with PixelLab (64x64 sprites). Same interface as sprite_parts.py.

The library is described by one JSON spec per base creature in pixellab_library/<body>.json
(schema in pixellab_library.py). The spec says how the art is generated (prompts, seeds),
how parts are extracted, and how they are picked (types, plans, tags). This module turns the
specs into the picking tables generate_sprite.py uses. Only parts that have been extracted
(parts_pl/<slot>/<name>.png exists) are registered.

Differences from the classic library:
- Every body is available to every type, but strongly prefers its own: a body's weight for a
  type is its `types` boost, or BODY_OTHER for types it doesn't list. Variants (edits of a
  body: same idea, different look) are bodies too and share their parent's plan and tags.
- Parts with a fixed-color ramp ("fixed": "fire", ...) are type-locked: they must list
  `types` and have other=0.
- `fits()` keeps every random pick inside the 64x64 canvas, so parts from any body can be
  offered to every body without hand-written size deny lists.
"""
import json
import sys
from pathlib import Path

from sprite_parts import MOTIF_PAIRS as CLASSIC_MOTIF_PAIRS
from sprite_parts import PLANS, TYPES, TYPE_COLORS, none, part  # noqa: F401  (TYPES/TYPE_COLORS re-exported)

NAME = "pixellab"
CANVAS = 64
GROUND_MARGIN = 3
PARTS_DIR = "parts_pl"
SPEC_DIR = Path(__file__).with_name("pixellab_library")

ALL = PLANS
FEATURE_SLOTS = ("head_top", "back", "tail")
DISABLED_SLOTS = {"aura"}
MAX_FEATURES = 2
MOTIF_SAME, MOTIF_CLASH = 2.5, 0.08
# Detailed generated art doesn't take painted-on spots well; markings stay off.
MARKINGS = {"none": {"_other": 1}}
# Weight of an empty slot, per slot. With MAX_FEATURES = 2 and head_top/back picked before tail,
# these give about: head_top 59%, back 60%, tail 51% filled, 3% plain (simulated over all types
# with 1208 bodies / 249 parts). Re-tune when the parts library grows a lot.
NONE_WEIGHT = {"head_top": 20.0, "back": 28.0, "tail": 1.5}
BODY_OTHER = 0.15  # a body's weight for types it doesn't list in `types` (listed types: 2-6)


class SpecError(Exception):
    pass


def load_specs(spec_dir=SPEC_DIR):
    """Every parseable spec; a file that doesn't parse (someone mid-edit) is skipped with a warning,
    and a spec with a top-level "skip": true (a rejected body) is left out entirely."""
    specs = []
    for path in sorted(Path(spec_dir).glob("*.json")):
        try:
            spec = json.loads(path.read_text())
        except (json.JSONDecodeError, OSError) as e:
            print(f"warning: skipping {path}: {e}", file=sys.stderr)
            continue
        if spec.get("skip"):
            continue
        spec["_path"] = str(path)
        specs.append(spec)
    return specs


def _plans(value):
    return list(ALL) if value in (None, "all") else list(value)


def _build(specs, parts_dir=PARTS_DIR):
    body, slots, motifs = [], {s: [] for s in FEATURE_SLOTS}, {}
    deny_from_parts = {}
    for s in specs:
        name = s["name"]
        for p in s.get("parts", []):
            if p.get("skip"):
                continue
            if p.get("fixed") and (p.get("other", 0) or not p.get("types")):
                print(f"warning: {s['_path']}: {p['name']} uses fixed colors '{p['fixed']}' but isn't type-locked "
                      f"(needs 'types' and other=0); left out", file=sys.stderr)
                continue
            if not (Path(parts_dir) / p["slot"] / f"{p['name']}.png").exists():
                continue
            palette = "fixed" if p.get("fixed") else p.get("palette", "main")
            entry = part(
                p["name"], _plans(p.get("plans")), p.get("types"), p.get("other", 0.3),
                tags=list(p.get("tags", [])) + [f"_slot_{p['slot']}"], requires=p.get("requires", ()),
                excludes=p.get("excludes", ()), palette=palette, z=p.get("z"),
                desc=p.get("desc", p.get("edit", "")) + f" (from {name})")
            entry["source"] = name          # the body it was generated on (see attached())
            slots[p["slot"]].append(entry)
            if p.get("motif"):
                motifs[p["name"]] = p["motif"]
            for b in p.get("deny_bodies", []):
                deny_from_parts.setdefault(b, []).append(p["name"])
        heads = [s["head"]] if isinstance(s.get("head"), str) else list(s.get("head", []))
        closed = [f"_slot_{slot}" for slot in FEATURE_SLOTS if slot not in s.get("slots", FEATURE_SLOTS)]
        # The body and its variants (edits of it: same idea, different look) share plan, head and tags.
        for b in [s] + [v for v in s.get("variants", []) if not v.get("skip")]:
            if not (Path(parts_dir) / "body" / f"{b['name']}.png").exists():
                continue
            entry = part(
                b["name"], [s["plan"]], b.get("types", s.get("types")), b.get("other", s.get("other", BODY_OTHER)),
                tags=heads + list(s.get("tags", [])), excludes=list(s.get("excludes", [])) + closed,
                deny=list(s.get("deny", [])), desc=b.get("desc", s.get("desc", "")))
            entry["family"] = name
            body.append(entry)
    for b in body:
        b["deny"] += deny_from_parts.get(b["name"], [])
    table = {"body": body, "aura": []}
    for slot in FEATURE_SLOTS:
        table[slot] = [none(other=NONE_WEIGHT[slot])] + slots[slot]
    return table, motifs


SLOTS, MOTIFS = _build(load_specs())
BODY = SLOTS["body"]
MOTIF_PAIRS = set(CLASSIC_MOTIF_PAIRS)

# Plan weight = total weight of that plan's bodies for the type, so picking a plan and then a body
# within it gives every body a chance proportional to its own weight.
PLAN_WEIGHTS = {t: {} for t in TYPES}
for _b in BODY:
    for _t in TYPES:
        _w = _b["types"].get(_t, _b["other"])
        PLAN_WEIGHTS[_t][_b["plans"][0]] = PLAN_WEIGHTS[_t].get(_b["plans"][0], 0) + _w


# ------------------------------------------------------------------ canvas fit + attachment

_extents = {}
_touch = {}


def _extent(slot, name):
    """(anchors, opaque bbox, opaque mask without marker pixels) of a compiled part, cached."""
    key = (slot, name)
    if key not in _extents:
        import numpy as np
        from PIL import Image

        import generate_sprite as gs   # lazy: generate_sprite imports this module by name
        img, anchors = gs.load_part(PARTS_DIR, slot, name)
        raw = np.array(Image.open(Path(PARTS_DIR) / slot / f"{name}.png").convert("RGBA"))
        mask = raw[..., 3] > 0
        for col in gs.MARKERS.values():
            mask &= ~np.all(raw[..., :3] == col, axis=-1)
        _extents[key] = (anchors, img.getbbox(), mask)
    return _extents[key]


MIN_COMPONENT = 6   # smaller specks are ignored by the attachment check


def _components(mask):
    """8-connected components of a mask with at least MIN_COMPONENT pixels, as (ys, xs) arrays."""
    import numpy as np
    seen = np.zeros_like(mask)
    h, w = mask.shape
    out = []
    for y0, x0 in zip(*np.nonzero(mask)):
        if seen[y0, x0]:
            continue
        seen[y0, x0] = True
        stack, comp = [(y0, x0)], []
        while stack:
            y, x = stack.pop()
            comp.append((y, x))
            for ny in (y - 1, y, y + 1):
                for nx in (x - 1, x, x + 1):
                    if 0 <= ny < h and 0 <= nx < w and mask[ny, nx] and not seen[ny, nx]:
                        seen[ny, nx] = True
                        stack.append((ny, nx))
        if len(comp) >= MIN_COMPONENT:
            out.append(tuple(np.array(comp).T))
    return out


_comps = {}


def contacts(body_name, slot, part_name):
    """For each sizeable piece of the part (placed at the body's anchor): does it touch the body
    (overlap or 8-adjacent)? Tuple of bools, one per piece, in a fixed order."""
    key = (body_name, slot, part_name)
    if key not in _touch:
        import numpy as np
        banch, _, bmask = _extent("body", body_name)
        panch, _, pmask = _extent(slot, part_name)
        if (slot, part_name) not in _comps:
            _comps[(slot, part_name)] = _components(pmask)
        grown = bmask.copy()
        for dy in (-1, 0, 1):
            for dx in (-1, 0, 1):
                grown |= np.roll(np.roll(np.pad(bmask, 1), dy, 0), dx, 1)[1:-1, 1:-1]
        ox, oy = banch[slot][0] - panch["pivot"][0], banch[slot][1] - panch["pivot"][1]
        hits = []
        for ys, xs in _comps[(slot, part_name)]:
            xs, ys = xs + ox, ys + oy
            inside = (xs >= 0) & (ys >= 0) & (xs < bmask.shape[1]) & (ys < bmask.shape[0])
            hits.append(bool(grown[ys[inside], xs[inside]].any()))
        _touch[key] = tuple(hits)
    return _touch[key]


_bodies = {b["name"] for b in BODY}


def attached(body, slot, p):
    """Every piece of the part that touches its own source body must touch this body too. A pair
    of wings needs both wings attached; a stemless tail or back spikes can't hang in mid-air;
    pieces that float free by design (halo, sparks, orbs) may float anywhere."""
    here = contacts(body["name"], slot, p["name"])
    source = p.get("source")
    home = contacts(source, slot, p["name"]) if source in _bodies else (True,) * len(here)
    return all(h or not s for h, s in zip(here, home))


def fits(body, picks):
    """True if `body` plus `picks` ({slot: part}) composes within CANVAS, using the same placement
    as generate_sprite.compose + place_on_canvas, and every part is attached to the body."""
    banch, bb, _ = _extent("body", body["name"])
    bx, by = banch["pivot"]
    boxes = [(bb[0] - bx, bb[1] - by, bb[2] - bx, bb[3] - by)]
    for slot, p in picks.items():
        if p["name"] == "none":
            continue
        if not attached(body, slot, p):
            return False
        anch, pb, _ = _extent(slot, p["name"])
        ax, ay = banch[slot]
        ox, oy = ax - bx - anch["pivot"][0], ay - by - anch["pivot"][1]
        boxes.append((pb[0] + ox, pb[1] + oy, pb[2] + ox, pb[3] + oy))
    w = max(b[2] for b in boxes) - min(b[0] for b in boxes)
    h = max(b[3] for b in boxes) - min(b[1] for b in boxes)
    return w + 2 <= CANVAS and h + GROUND_MARGIN + 1 <= CANVAS