#!/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 build_part(name, variant, out, subdiv=48):
    """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(ba._head_proto(mats, spec, subdiv=subdiv), "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, head_proto=False):
    """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. Pass
    ``head_proto=True`` to build with the connected-shell prototype head.
    """
    objs = ba.build_character(variant, head_proto=head_proto)
    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)
        s.add_argument("--subdiv", type=int, default=48,
                       help="head_proto UV-sphere u-segments (24/48/64 variants)")

    asm = sub.add_parser("assemble")
    asm.add_argument("--variant", default=PROTOTYPE)
    asm.add_argument("--out", required=True)
    asm.add_argument("--head-proto", action="store_true",
                     help="use the connected-shell prototype head")

    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, subdiv=args.subdiv)
        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, head_proto=args.head_proto)
        print(f"assembled {args.variant} -> {args.out}")
    elif args.cmd == "inspect":
        inspect(args.file)
    return 0


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