"""
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 (`other` defaults to 1); `types` only boosts.
- 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
from pathlib import Path

from sprite_parts import MOTIF_PAIRS as CLASSIC_MOTIF_PAIRS
from sprite_parts import PLAN_WEIGHTS as CLASSIC_PLAN_WEIGHTS
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 69%, back 55%, tail 51% filled, 2% plain (simulated over all types
# with 31 bodies / 111 parts). Re-tune when the library grows a lot.
NONE_WEIGHT = {"head_top": 10.0, "back": 16.0, "tail": 1.0}
PLAN_FLOOR = 1.0   # every plan stays possible for every type


class SpecError(Exception):
    pass


def load_specs(spec_dir=SPEC_DIR):
    specs = []
    for path in sorted(Path(spec_dir).glob("*.json")):
        spec = json.loads(path.read_text())
        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")):
                raise SpecError(f"{s['_path']}: {p['name']} uses fixed colors '{p['fixed']}' and must be "
                                f"type-locked: give it 'types' and other=0")
            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"])
        if not (Path(parts_dir) / "body" / f"{name}.png").exists():
            continue
        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)]
        body.append(part(
            name, [s["plan"]], s.get("types"), s.get("other", 1.0), tags=heads + list(s.get("tags", [])),
            excludes=list(s.get("excludes", [])) + closed, deny=list(s.get("deny", [])), desc=s.get("desc", "")))
    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)

_present = {b["plans"][0] for b in BODY}
PLAN_WEIGHTS = {t: {plan: CLASSIC_PLAN_WEIGHTS[t].get(plan, 0) + PLAN_FLOOR for plan in PLANS if plan in _present}
                for t in TYPES}


# ------------------------------------------------------------------ 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