#!/usr/bin/env python3
"""Local part pipeline for the Blender NPC model overhaul.

Builds, extracts, assembles, inspects and sheets the reusable head/face and hand
parts before they are rolled into the shipped characters. It reuses the
builders in ``blender_assets.py`` (no second face/hand implementation) and stays
local-only and headless: the Blender commands run inside the flatpak, and the
sheet/compare commands run with stdlib + Pillow.

Part sources and preview GLBs live under ``tmp/npc-parts/<variant>/`` (gitignored).
Only the assembled character GLBs, AO/maps, the manifest and the review sheets
are committed; a part GLB is not a shipped asset.

    # Build one part fresh from the current builders (head|face|hand_l|hand_r)
    flatpak run org.blender.Blender --background \
      --python "$PWD/tools/npc_parts.py" -- build --part face \
      --variant npc_casual --out "$PWD/tmp/npc-parts/npc_casual/face.glb"
    # Pull the current shipped head/hands out of the built character GLB
    flatpak run org.blender.Blender --background \
      --python "$PWD/tools/npc_parts.py" -- extract --part hand_l \
      --variant npc_casual --out "$PWD/tmp/npc-parts/npc_casual/hand_l.glb"
    # Assemble a variant (currently rebuilds the procedural character)
    flatpak run org.blender.Blender --background \
      --python "$PWD/tools/npc_parts.py" -- assemble --variant npc_casual \
      --out "$PWD/tmp/npc-parts/npc_casual/assembled.glb"
    # Inspect a part or assembled GLB (Blender import view)
    flatpak run org.blender.Blender --background \
      --python "$PWD/tools/npc_parts.py" -- inspect --file "$PWD/tmp/npc-parts/npc_casual/face.glb"
    # Contact sheet / before-after comparison (outside Blender)
    python3 tools/npc_parts.py sheet --dir tmp/npc-parts/npc_casual --out docs/blender_previews/sheet_parts.jpg
    python3 tools/npc_parts.py compare --before DIR --after DIR --out sheet.jpg

The flatpak sandbox cwd is /app/blender and cannot read /tmp, so every path is
absolute and derived from this file's location.
"""

import argparse
import math
import sys
from pathlib import Path

# The flatpak sandbox may not put this script's directory on sys.path, so add it
# explicitly so `import blender_assets` works from any cwd.
_TOOLS = str(Path(__file__).resolve().parent)
if _TOOLS not in sys.path:
    sys.path.insert(0, _TOOLS)

import blender_assets as ba  # same directory; no second implementation

ROOT = Path(__file__).resolve().parent.parent
PARTS_DIR = ROOT / "tmp" / "npc-parts"
PROTOTYPE = "npc_casual"  # canonical one-variant prototype

# The four buildable parts. `head` is the bare head shell; `face` is the shell
# plus nose/eyes/mouth/ears; `hand_l`/`hand_r` are the left/right hand geometry.
PART_NAMES = ("head", "face", "hand_l", "hand_r", "head_proto")
_HAND_PREFIXES = ("Palm", "Finger", "Fingertip", "Thumb", "ThumbTip")


def _needs_blender(cmd):
    if ba.bpy is not None:
        return None
    print(
        f"npc_parts {cmd} runs inside Blender. Use:\n"
        f"  flatpak run org.blender.Blender --background \\\n"
        f"    --python \"$PWD/tools/npc_parts.py\" -- {cmd} ...",
        file=sys.stderr,
    )
    return 2


def _mats_and_spec(variant):
    spec = dict(ba.CHARACTERS[variant])
    mats = {slot: ba.material(slot) for slot in ba.SLOTS}
    for extra in ("hivis", "reflective", "leather",
                  "steel", "patch", "plastic_black", "rubber", "helmet"):
        mats[extra] = ba.material(extra, textured_solid=True)
    return mats, spec


def _export_obj(obj, path):
    """Export one object as a selection-only GLB (no .blend source side-effect)."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    for o in ba.bpy.context.selected_objects:
        o.select_set(False)
    obj.select_set(True)
    ba.bpy.context.view_layer.objects.active = obj
    ba.bpy.ops.export_scene.gltf(
        filepath=str(path), export_format="GLB", use_selection=True,
        export_apply=True, export_animations=False, export_skins=True,
        export_yup=True, export_materials="EXPORT", export_image_format="AUTO")
    return path


def _head_proto(mats, spec):
    """Connected head shell with sculpted sockets, nose, lips, ears and brow/cheek/jaw.

    This is the §3 prototype: the skin features are displaced out of one head
    mesh instead of floating spheres, so the face is a single connected surface
    with loop-based eye sockets and protruding nose/lip/ear volumes. The eyeballs
    (sclera/iris/lids) stay separate — they are not skin. Blender Z-up, face
    front −Y.
    """
    skin, acc, hair = mats["skin"], mats["accessory"], mats["hair"]
    white = ba.material("eye_white", textured_solid=True)
    head = ba.sphere("HeadShell", 0.155, (0, 0, 1.635), skin, seg=48, rings=26,
                     scale=(0.90, 0.96, 1.08))
    ba._shape_head_mesh(head)  # brow ridge + cheekbone + jaw/chin planes
    bm = ba.bmesh.new()
    bm.from_mesh(head.data)

    # Pass 1: skull shape — flatten the face plane and taper the jaw/chin so the
    # head stops reading as an egg.
    for v in bm.verts:
        x, y, z = v.co
        if y < -0.03 and 1.44 < z < 1.76:
            v.co.y = -0.13 + (y + 0.13) * 0.10  # flatten the forward bulge
        if y < -0.03 and z > 1.70:
            v.co.y += (z - 1.70) * 0.35  # forehead slopes back to the crown
        if z < 1.58:
            taper = max(0.0, 1.0 - (1.58 - z) / 0.12)
            v.co.x *= 1.0 - 0.16 * taper  # narrow the lower face

    # Pass 2: features on the flattened face.
    for v in bm.verts:
        x, y, z = v.co
        ax = abs(x)
        # --- brow step: a clear ridge where the forehead meets the nose ---
        brow = (max(0.0, 1.0 - abs(z - 1.695) / 0.022)
                * max(0.0, 1.0 - abs(ax - 0.050) / 0.045))
        if brow > 0:
            v.co.y -= 0.016 * brow
        # --- eye sockets: recess the floor, leave a raised orbital rim ---
        for s in (-1.0, 1.0):
            r = math.hypot(x - 0.052 * s, z - 1.664)
            if r < 0.036:
                v.co.y += 0.018 * (1.0 - r / 0.036)  # recess inward (+y)
            elif r < 0.046:
                v.co.y -= 0.010 * (1.0 - (r - 0.036) / 0.010)  # rim forward
        # --- nose: a wide bridge/tip with ala wings, not a spike ---
        if y < 0 and ax < 0.075:
            band = max(0.0, 1.0 - abs(z - 1.632) / 0.055)  # 1.577..1.687
            if band > 0:
                tip = max(0.0, 1.0 - abs(z - 1.615) / 0.032)
                v.co.y -= (0.060 * tip + 0.018 * max(0.0, band - tip)) \
                          * (1.0 - ax / 0.075)
            for s in (-1.0, 1.0):
                ar = math.hypot((x - 0.024 * s) / 0.014, (z - 1.604) / 0.013)
                if ar < 1.0:
                    v.co.y -= 0.022 * (1.0 - ar)
        # --- chin: a distinct forward bump below the mouth ---
        if y < 0 and ax < 0.040:
            chin = max(0.0, 1.0 - abs(z - 1.548) / 0.022)
            if chin > 0:
                v.co.y -= 0.030 * chin
            # jaw step: recess just below the chin so it reads as a chin-to-neck
            step = max(0.0, 1.0 - abs(z - 1.516) / 0.016)
            if step > 0:
                v.co.y += 0.024 * step
        # --- mouth: recessed line + upper/lower lip volumes ---
        if y < 0 and ax < 0.030:
            if abs(z - 1.583) < 0.012:
                v.co.y += 0.008  # mouth slit recess
            if abs(z - 1.590) < 0.008:
                v.co.y -= 0.014  # upper lip
            if abs(z - 1.577) < 0.010:
                v.co.y -= 0.018  # fuller lower lip
    # --- ears: helix rim + concha recess + a lobe ---
    for v in bm.verts:
        x, y, z = v.co
        r = math.hypot((z - 1.60) / 0.055, (y - 0.002) / 0.045)
        near = abs(x) - 0.138
        sign = 1.0 if x >= 0 else -1.0
        if -0.02 < near < 0.06 and r < 1.0:
            if r < 0.30:
                v.co.x -= 0.018 * (1.0 - r / 0.30) * sign  # concha recess (inward)
            else:
                v.co.x += 0.045 * (1.0 - r) * sign  # helix rim outward
        # lobe: a distinct bump at the lower front of the ear
        lr = math.hypot((z - 1.575) / 0.016, (y + 0.012) / 0.016)
        if -0.02 < near < 0.05 and lr < 1.0:
            v.co.x += 0.032 * (1.0 - lr) * sign
    bm.normal_update()
    bm.to_mesh(head.data)
    bm.free()

    # Eyeballs stay separate for readability; lids wrap the socket rim.
    parts = [head]
    for s in (-1.0, 1.0):
        ex, ey, ez = 0.052 * s, -0.134, 1.664
        parts.append(ba.sphere("EyeWhite", 0.019, (ex, ey, ez), white, seg=10,
                               rings=6, scale=(1.25, 0.62, 0.85)))
        parts.append(ba.sphere("Iris", 0.0105, (ex, ey - 0.007, ez), acc, seg=8,
                               rings=5, scale=(1.0, 0.55, 1.0)))
        parts.append(ba.sphere("UpperLid", 0.020, (ex, ey - 0.006, ez + 0.014),
                               skin, seg=8, rings=4, scale=(1.20, 0.32, 0.24),
                               rot=(math.radians(-10), 0, 0)))
        parts.append(ba.sphere("LowerLid", 0.017, (ex, ey - 0.004, ez - 0.012),
                               skin, seg=8, rings=4, scale=(1.10, 0.22, 0.18),
                               rot=(math.radians(6), 0, 0)))
        parts.append(ba.box("Brow", (0.055, 0.014, 0.012), (0.052 * s, -0.143, 1.699),
                            hair, rot=(0, 0.14 * s, 0.0)))
    return parts


def build_part(name, variant, out):
    """Build one part fresh from the current builders (Blender-only)."""
    mats, spec = _mats_and_spec(variant)
    if name == "head":
        obj = ba.sphere("Head", 0.155, (0, 0, 1.635), mats["skin"], seg=24,
                        rings=16, scale=(0.92, 1.0, 1.12))
        ba._shape_head_mesh(obj)
    elif name == "face":
        obj = ba.join(ba._face(mats, spec), "face")
    elif name == "head_proto":
        obj = ba.join(_head_proto(mats, spec), "head_proto")
    else:
        side = -1.0 if name == "hand_l" else 1.0
        hs = ba._hands_and_shoes(mats, spec)

        def _center_x(obj):
            xs = [v.co.x for v in obj.data.vertices]
            return (min(xs) + max(xs)) / 2

        hand = [o for o in hs if o.name.startswith(_HAND_PREFIXES)
                and (_center_x(o) < 0 if side < 0 else _center_x(o) > 0)]
        obj = ba.join(hand, name)
    return _export_obj(obj, out)


def _extract_group(mesh, group_name, out):
    """Copy the vertices of one vertex group into a new standalone mesh."""
    vg = mesh.vertex_groups.get(group_name)
    if vg is None:
        raise SystemExit(f"mesh {mesh.name} has no vertex group {group_name!r}")
    bm = ba.bmesh.new()
    bm.from_mesh(mesh.data)
    deform = bm.verts.layers.deform.active
    keep = {v.index for v in bm.verts
            if deform is not None and vg.index in v[deform]}
    if not keep:
        bm.free()
        raise SystemExit(f"vertex group {group_name!r} is empty in {mesh.name}")
    doomed = [v for v in bm.verts if v.index not in keep]
    ba.bmesh.ops.delete(bm, geom=doomed, context="VERTS")
    bm.to_mesh(mesh.data)
    bm.free()
    return _export_obj(mesh, out)


def extract_part(name, variant, out):
    """Pull the current shipped head/hands out of the built character GLB.

    The shipped mesh joins the head shell and its features into one Head-weighted
    piece, so ``head`` and ``face`` both extract that piece; ``hand_l``/``hand_r``
    extract the HandL/HandR-weighted geometry.
    """
    src = ba.asset_path("npc", variant)
    if not src.is_file():
        raise SystemExit(f"not built: {src}; run blender_assets.py build first")
    ba.bpy.ops.wm.read_factory_settings(use_empty=True)
    ba.bpy.ops.import_scene.gltf(filepath=str(src))
    mesh = next((o for o in ba.bpy.context.scene.objects if o.type == "MESH"
                 and o.name == "npc_base"), None)
    if mesh is None:
        raise SystemExit(f"no npc_base mesh in {src}")
    group = ("Head" if name in ("head", "face", "head_proto")
             else f"Hand{'L' if name == 'hand_l' else 'R'}")
    return _extract_group(mesh, group, out)


def assemble(variant, out):
    """Assemble a variant (currently rebuilds the procedural character).

    This proves the end-to-end wiring; once parts are approved, this swaps the
    new part GLBs in instead of calling ``build_character`` directly.
    """
    objs = ba.build_character(variant)
    mesh = next(o for o in objs if o.type == "MESH")
    arm = next(o for o in objs if o.type == "ARMATURE")
    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    for o in ba.bpy.context.selected_objects:
        o.select_set(False)
    for o in (arm, mesh):
        o.select_set(True)
    ba.bpy.context.view_layer.objects.active = mesh
    ba.bpy.ops.export_scene.gltf(
        filepath=str(out), export_format="GLB", use_selection=True,
        export_apply=True, export_animations=False, export_skins=True,
        export_yup=True, export_materials="EXPORT", export_image_format="AUTO")
    return out


def inspect(path):
    result = ba.inspect_glb(str(path), category="npc")
    ba.print_inspection(result)
    return result


def _pillow():
    try:
        from PIL import Image, ImageOps  # noqa: F401
        return Image, ImageOps
    except ImportError:
        raise SystemExit("Pillow is required for sheet/compare (pip install pillow)")


def _pngs(directory):
    d = Path(directory)
    if not d.is_dir():
        raise SystemExit(f"not a directory: {d}")
    files = sorted(d.glob("*.png"))
    if not files:
        raise SystemExit(f"no PNGs in {d}")
    return files


def make_sheet(directory, out, cols=4):
    Image, ImageOps = _pillow()
    files = _pngs(directory)
    thumb_w = 512
    tiles = []
    for f in files:
        im = Image.open(f).convert("RGB")
        im = ImageOps.fit(im, (thumb_w, thumb_w), Image.LANCZOS)
        tiles.append(im)
    rows = (len(tiles) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * thumb_w, rows * thumb_w), (24, 24, 24))
    for i, im in enumerate(tiles):
        sheet.paste(im, ((i % cols) * thumb_w, (i // cols) * thumb_w))
    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(out, quality=90)
    print(f"sheet: {out} ({len(tiles)} tiles)")
    return out


def compare(before, after, out, cols=2):
    Image, ImageOps = _pillow()
    before_files = _pngs(before)
    after_files = _pngs(after)
    if len(before_files) != len(after_files):
        print(f"warn: {len(before_files)} before vs {len(after_files)} after PNGs",
              file=sys.stderr)
    pairs = max(len(before_files), len(after_files))
    thumb_w = 512
    tiles = []
    for i in range(pairs):
        for files in (before_files, after_files):
            if i < len(files):
                im = Image.open(files[i]).convert("RGB")
                im = ImageOps.fit(im, (thumb_w, thumb_w), Image.LANCZOS)
            else:
                im = Image.new("RGB", (thumb_w, thumb_w), (24, 24, 24))
            tiles.append(im)
    rows = (len(tiles) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * thumb_w, rows * thumb_w), (24, 24, 24))
    for i, im in enumerate(tiles):
        sheet.paste(im, ((i % cols) * thumb_w, (i // cols) * thumb_w))
    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(out, quality=90)
    print(f"compare: {out} ({pairs} pairs)")
    return out


def main(argv=None):
    raw = list(argv if argv is not None else sys.argv[1:])
    if "--" in raw:
        raw = raw[raw.index("--") + 1:]
    ap = argparse.ArgumentParser(description="Local NPC part pipeline")
    sub = ap.add_subparsers(dest="cmd", required=True)

    for cmd in ("build", "extract"):
        s = sub.add_parser(cmd)
        s.add_argument("--part", required=True, choices=PART_NAMES)
        s.add_argument("--variant", default=PROTOTYPE)
        s.add_argument("--out", required=True)

    asm = sub.add_parser("assemble")
    asm.add_argument("--variant", default=PROTOTYPE)
    asm.add_argument("--out", required=True)

    insp = sub.add_parser("inspect")
    insp.add_argument("--file", required=True)

    sh = sub.add_parser("sheet")
    sh.add_argument("--dir", required=True)
    sh.add_argument("--out", required=True)
    sh.add_argument("--cols", type=int, default=4)

    cm = sub.add_parser("compare")
    cm.add_argument("--before", required=True)
    cm.add_argument("--after", required=True)
    cm.add_argument("--out", required=True)

    args = ap.parse_args(raw)

    if args.cmd == "sheet":
        make_sheet(args.dir, args.out, args.cols)
        return 0
    if args.cmd == "compare":
        compare(args.before, args.after, args.out)
        return 0

    # Everything else runs inside Blender.
    code = _needs_blender(args.cmd)
    if code is not None:
        return code
    ba.reset_scene()
    if args.cmd == "build":
        build_part(args.part, args.variant, args.out)
        print(f"built part {args.part} -> {args.out}")
    elif args.cmd == "extract":
        extract_part(args.part, args.variant, args.out)
        print(f"extracted {args.part} -> {args.out}")
    elif args.cmd == "assemble":
        assemble(args.variant, args.out)
        print(f"assembled {args.variant} -> {args.out}")
    elif args.cmd == "inspect":
        inspect(args.file)
    return 0


if __name__ == "__main__":
    sys.exit(main())