"""
blender_author_offset_shell.py  (T-1089, route (c) offset-shell authoring)

Derive a garment SHELL from our own body segment meshes — the "create-stuff-
yourself" authoring pipeline that owes nothing to any vendor pack. Because the
shell IS our body topology, bone weights are inherited by construction (no
Surface-Deform, no Data-Transfer, no re-rig): every vertex keeps the 65-bone
vertex groups it had as skin.

Pipeline (t-shirt reference on average_m):
  1. Import the body segments the garment covers (torso + torso_upper + upper
     arms), keep only the skinned body meshes (Icosphere debris filtered out).
  2. Join into one mesh under a single armature; merge vertex groups by name.
  3. Bone-plane SLEEVE cut — trim the upper-arm tube to short-sleeve length via
     a coordinate threshold derived from the upperarm bone axis (robust; no
     boundary-loop classification, which the segment tool warns is fragile).
     Neckline + hem come free as the natural segment boundaries.
  4. Offset the surface outward along vertex normals (~12 mm standoff from skin).
  5. Solidify (use_rim=True) — gives the cloth real thickness and caps the cut
     rims (sleeve openings) into hems.
  6. Assign ONE flat modern-fabric material (drops all skin textures). Style pin:
     neutral heather tone, no trim, no fantasy anything.
  7. Bake a UV0-aligned RGBA REGION MASK per-face: collar band -> R, main body
     -> G, sleeve trim -> B (channel-routed 4-tint shader input, G3).
  8. Author a 2nd UV channel (TEXCOORD_1) projecting the front chest into [0,1]
     for the logo decal (G4); everything else parks outside the box.
  9. Export the reference GLB (export_skins=True) + write reference_mask.png and
     base_albedo.png sidecars.

The output reference is authored on average_m only; G1
(blender_batch_fit_skinned.py) fits it to the other body types.

Run:
  tooling/blender --background --python \
    tooling/garment-fit/blender_author_offset_shell.py -- \
    <bodies_dir>/average_m <out_dir> [--offset 0.012] [--sleeve-frac 0.40]

Writes:
  <out_dir>/average_m.glb       reference garment (skinned)
  <out_dir>/reference_mask.png  RGBA region mask (UV0-aligned)
  <out_dir>/base_albedo.png     flat fabric albedo (also embedded in the GLB)

Decisions: D-162 (clothing pre-fitted per body type), D-251 (in-house wardrobe).
"""

import sys
import os
import bpy
import bmesh
import numpy as np
from mathutils import Vector

# --------------------------------------------------------------------------
# Parameters
# --------------------------------------------------------------------------

# Segments a t-shirt covers. Order matters only for join-active choice.
COVERED_SEGMENTS = ["seg_torso", "seg_torso_upper", "seg_arm_upper_l", "seg_arm_upper_r"]

OFFSET_M = 0.020          # outward standoff along vertex normals; larger than the
                          # 10-14 mm ideal on the reference body buys clearance for
                          # bigger bodies under Surface-Deform batch-fit (Q-060) —
                          # the muscular/female torso otherwise pokes through.
CLOTH_THICKNESS_M = 0.004  # Solidify thickness after offset (total ~24 mm standoff)
SLEEVE_FRAC = 0.40        # fraction of upper-arm length kept (short sleeve)
MASK_SIZE = 512
ALBEDO_SIZE = 512

# Flat modern-fabric base tone (linear-ish sRGB), neutral heather grey.
FABRIC_RGB = (0.60, 0.61, 0.63)
FABRIC_NOISE = 0.03       # +/- albedo jitter for a subtle woven feel

# Chest logo box in body-local metres (X width, Z height).
# FRONT AXIS: these Quaternius bodies face -Y in Blender (verified empirically —
# a +Y test projection landed on the character's back). So "front" = -Y.
FRONT_Y_SIGN = -1.0
CHEST_X = (-0.12, 0.12)
CHEST_Z = (1.16, 1.44)
CHEST_FRONT_Y = 0.015     # face centre must be on the front side by at least this
CHEST_NORMAL_Y = 0.20     # face normal must point forward by at least this much

# Region-mask classification (body-local, Z up).
COLLAR_Z_MIN = 1.49       # faces above this AND near centre -> collar band (R)
COLLAR_X_ABS = 0.11       # collar band stays near the neck, not the shoulders
SLEEVE_X_ABS = 0.20       # faces with |center X| beyond this -> sleeve cap (B)


def log(msg):
    print(f"[offset-shell] {msg}")


# --------------------------------------------------------------------------
# Scene helpers
# --------------------------------------------------------------------------

def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    bpy.ops.outliner.orphans_purge(do_recursive=True)


def import_glb(path):
    before = set(bpy.context.scene.objects)
    bpy.ops.import_scene.gltf(filepath=path)
    return [o for o in bpy.context.scene.objects if o not in before]


def is_body_mesh(obj):
    """A real skinned body segment mesh — not Icosphere debris."""
    if obj.type != 'MESH':
        return False
    if obj.name.startswith("Icosphere"):
        return False
    if len(obj.vertex_groups) == 0:
        return False
    if len(obj.data.vertices) < 50:
        return False
    return True


# --------------------------------------------------------------------------
# Build the joined shell base
# --------------------------------------------------------------------------

def build_covered_mesh(body_dir):
    """Import covered segments, keep skinned meshes, join to one mesh + armature."""
    body_meshes = []
    armature = None

    for seg in COVERED_SEGMENTS:
        path = os.path.join(body_dir, f"{seg}.glb")
        if not os.path.isfile(path):
            log(f"WARNING: missing segment {path} — skipping")
            continue
        objs = import_glb(path)
        for o in objs:
            if o.type == 'ARMATURE' and armature is None:
                armature = o
            elif o.type == 'ARMATURE':
                # drop extra armature copies (identical rest pose)
                bpy.data.objects.remove(o, do_unlink=True)
            elif is_body_mesh(o):
                body_meshes.append(o)
            else:
                # Icosphere / debris
                bpy.data.objects.remove(o, do_unlink=True)

    if not body_meshes:
        raise RuntimeError("no skinned body meshes imported for covered segments")
    if armature is None:
        raise RuntimeError("no armature found in covered segments")

    # Join meshes (vertex groups merge by name across segments).
    bpy.ops.object.select_all(action='DESELECT')
    for m in body_meshes:
        m.select_set(True)
    bpy.context.view_layer.objects.active = body_meshes[0]
    bpy.ops.object.join()
    shell = bpy.context.active_object
    shell.name = "garment_shell"

    # Re-point the armature modifier at the surviving armature; re-parent.
    for mod in list(shell.modifiers):
        if mod.type == 'ARMATURE':
            mod.object = armature
    shell.parent = armature
    shell.matrix_parent_inverse = armature.matrix_world.inverted()

    log(f"joined shell: {len(shell.data.vertices)} verts, "
        f"{len(shell.data.polygons)} faces, {len(shell.vertex_groups)} vgroups")
    return shell, armature


# --------------------------------------------------------------------------
# Bone-plane sleeve cut
# --------------------------------------------------------------------------

def sleeve_cut(shell, armature):
    """Delete sleeve-tip verts beyond the short-sleeve plane on each upper arm.

    The upper arm runs along +/-X (shoulder head -> elbow tail). We keep the
    fraction SLEEVE_FRAC of that length from the shoulder and delete the rest.
    Torso verts stay (|X| < shoulder head), so a single coordinate threshold is
    safe and needs no per-vertex weight test.
    """
    bones = armature.data.bones
    cut_planes = []  # (axis_sign, threshold_x)
    for bone_name, sign in [("upperarm_l", +1), ("upperarm_r", -1)]:
        b = bones.get(bone_name)
        if b is None:
            log(f"WARNING: bone {bone_name} missing — sleeve not cut on that side")
            continue
        head_x = b.head_local.x
        tail_x = b.tail_local.x
        thr = head_x + SLEEVE_FRAC * (tail_x - head_x)
        cut_planes.append((sign, thr))
        log(f"sleeve cut {bone_name}: keep |x| up to {thr:.3f} "
            f"(shoulder {head_x:.3f} -> elbow {tail_x:.3f})")

    bm = bmesh.new()
    bm.from_mesh(shell.data)
    bm.verts.ensure_lookup_table()
    to_delete = []
    for v in bm.verts:
        for sign, thr in cut_planes:
            if sign > 0 and v.co.x > thr:
                to_delete.append(v)
                break
            if sign < 0 and v.co.x < thr:
                to_delete.append(v)
                break
    bmesh.ops.delete(bm, geom=to_delete, context='VERTS')
    bm.to_mesh(shell.data)
    bm.free()
    shell.data.update()
    log(f"sleeve cut removed {len(to_delete)} verts; "
        f"{len(shell.data.vertices)} remain")


# --------------------------------------------------------------------------
# Outward offset + solidify
# --------------------------------------------------------------------------

def offset_outward(shell, offset):
    """Push every vertex outward along its (smoothed) normal by `offset` m."""
    me = shell.data
    me.calc_normals_split() if hasattr(me, "calc_normals_split") else None
    bm = bmesh.new()
    bm.from_mesh(me)
    bm.normal_update()
    for v in bm.verts:
        v.co += v.normal * offset
    bm.to_mesh(me)
    bm.free()
    me.update()
    log(f"offset surface outward by {offset*1000:.0f} mm along normals")


def solidify(shell, thickness):
    """Solidify with use_rim to give cloth thickness and cap the cut rims."""
    bpy.ops.object.select_all(action='DESELECT')
    shell.select_set(True)
    bpy.context.view_layer.objects.active = shell
    # consistent outward normals first
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.object.mode_set(mode='OBJECT')
    sol = shell.modifiers.new(name="Solidify", type='SOLIDIFY')
    sol.thickness = thickness
    sol.offset = 1.0        # grow outward only
    sol.use_rim = True      # cap open boundaries (sleeve/neck/hem)
    sol.use_rim_only = False
    bpy.ops.object.modifier_apply(modifier=sol.name)
    log(f"solidified: {thickness*1000:.0f} mm, use_rim; "
        f"{len(shell.data.vertices)} verts")


# --------------------------------------------------------------------------
# Material (flat fabric albedo)
# --------------------------------------------------------------------------

def make_base_albedo_image(seed=1089):
    img = bpy.data.images.new("garment_base_albedo", ALBEDO_SIZE, ALBEDO_SIZE, alpha=False)
    rng = np.random.default_rng(seed)
    base = np.array(FABRIC_RGB, dtype=np.float32)
    noise = (rng.random((ALBEDO_SIZE * ALBEDO_SIZE, 1), dtype=np.float32) - 0.5) * 2.0 * FABRIC_NOISE
    rgb = np.clip(base[None, :] + noise, 0.0, 1.0)
    rgba = np.concatenate([rgb, np.ones((ALBEDO_SIZE * ALBEDO_SIZE, 1), dtype=np.float32)], axis=1)
    img.pixels.foreach_set(rgba.reshape(-1))
    img.update()
    return img


def assign_fabric_material(shell, albedo_img):
    shell.data.materials.clear()
    mat = bpy.data.materials.new("garment_fabric")
    mat.use_nodes = True
    nt = mat.node_tree
    bsdf = nt.nodes.get("Principled BSDF")
    tex = nt.nodes.new("ShaderNodeTexImage")
    tex.image = albedo_img
    nt.links.new(tex.outputs["Color"], bsdf.inputs["Base Color"])
    if "Roughness" in bsdf.inputs:
        bsdf.inputs["Roughness"].default_value = 0.9
    shell.data.materials.append(mat)
    log("assigned flat fabric material (skin textures dropped)")


# --------------------------------------------------------------------------
# Region mask bake (UV0-aligned, per-face rasterization)
# --------------------------------------------------------------------------

def _classify_region(center):
    """Return an RGBA region colour for a face centre (body-local coords)."""
    x, y, z = center.x, center.y, center.z
    if abs(x) >= SLEEVE_X_ABS:
        return (0.0, 0.0, 1.0, 0.0)   # sleeve caps -> B (tint[2])
    if z >= COLLAR_Z_MIN and abs(x) < COLLAR_X_ABS:
        return (1.0, 0.0, 0.0, 0.0)   # neck collar band -> R (tint[0])
    return (0.0, 1.0, 0.0, 0.0)       # main body -> G (tint[1])


def _tris_from_face(face, uv_layer):
    """Fan-triangulate a bmesh face into (uv, uv, uv) tuples in [0,1] space."""
    loops = face.loops[:]
    uvs = [l[uv_layer].uv.copy() for l in loops]
    tris = []
    for i in range(1, len(uvs) - 1):
        tris.append((uvs[0], uvs[i], uvs[i + 1]))
    return tris


def bake_region_mask(shell, out_path):
    """Rasterize each face's UV0 triangle with its region colour into MASK_SIZE^2.

    Background initialised to main-body green so bilinear bleed at island edges
    never lands on an untinted (all-zero) texel.
    """
    W = H = MASK_SIZE
    buf = np.zeros((H, W, 4), dtype=np.float32)
    buf[:, :, 1] = 1.0  # green background = main body

    me = shell.data
    bm = bmesh.new()
    bm.from_mesh(me)
    bm.faces.ensure_lookup_table()
    uv_layer = bm.loops.layers.uv.active
    if uv_layer is None:
        raise RuntimeError("no active UV layer for region mask bake")

    yy, xx = np.mgrid[0:H, 0:W]
    for face in bm.faces:
        color = _classify_region(face.calc_center_median())
        for a, b, c in _tris_from_face(face, uv_layer):
            _raster_tri(buf, a, b, c, color, W, H)
    bm.free()

    # Blender image is bottom-up; buf row 0 is V=0 (bottom) already since we
    # rasterise with row = v*(H-1). Save via Blender to match the texture pipe.
    img = bpy.data.images.new("garment_region_mask", W, H, alpha=True)
    img.pixels.foreach_set(buf.reshape(-1))
    img.update()
    img.filepath_raw = out_path
    img.file_format = 'PNG'
    img.save()
    log(f"baked region mask -> {out_path}")


def _raster_tri(buf, a, b, c, color, W, H):
    """Barycentric fill of a UV triangle into buf (V=0 at row 0 = bottom)."""
    ax, ay = a.x * (W - 1), a.y * (H - 1)
    bx, by = b.x * (W - 1), b.y * (H - 1)
    cx, cy = c.x * (W - 1), c.y * (H - 1)
    minx = max(int(np.floor(min(ax, bx, cx))), 0)
    maxx = min(int(np.ceil(max(ax, bx, cx))), W - 1)
    miny = max(int(np.floor(min(ay, by, cy))), 0)
    maxy = min(int(np.ceil(max(ay, by, cy))), H - 1)
    if minx > maxx or miny > maxy:
        return
    denom = (by - cy) * (ax - cx) + (cx - bx) * (ay - cy)
    if abs(denom) < 1e-9:
        return
    ys, xs = np.mgrid[miny:maxy + 1, minx:maxx + 1]
    px = xs + 0.5
    py = ys + 0.5
    w0 = ((by - cy) * (px - cx) + (cx - bx) * (py - cy)) / denom
    w1 = ((cy - ay) * (px - cx) + (ax - cx) * (py - cy)) / denom
    w2 = 1.0 - w0 - w1
    inside = (w0 >= -1e-4) & (w1 >= -1e-4) & (w2 >= -1e-4)
    if not inside.any():
        return
    region = buf[miny:maxy + 1, minx:maxx + 1, :]
    col = np.array(color, dtype=np.float32)
    region[inside] = col


# --------------------------------------------------------------------------
# Logo UV2 chest channel
# --------------------------------------------------------------------------

def author_logo_uv(shell):
    """Create a 2nd UV layer projecting front chest faces into [0,1]; park the
    rest outside the box (shader guards uv2 in [0,1])."""
    me = shell.data
    # Keep exactly two UV layers: primary (albedo/mask) + logo. Remove extras.
    while len(me.uv_layers) > 1:
        me.uv_layers.remove(me.uv_layers[-1])
    logo_uv = me.uv_layers.new(name="logo_uv")
    me.uv_layers.active = me.uv_layers[0]  # keep albedo layer active for mask bake safety

    bm = bmesh.new()
    bm.from_mesh(me)
    bm.faces.ensure_lookup_table()
    bm.normal_update()
    uvl = bm.loops.layers.uv.get("logo_uv")

    x0, x1 = CHEST_X
    z0, z1 = CHEST_Z
    placed = 0
    for face in bm.faces:
        center = face.calc_center_median()
        on_chest = (
            center.y * FRONT_Y_SIGN > CHEST_FRONT_Y
            and x0 <= center.x <= x1
            and z0 <= center.z <= z1
            and face.normal.y * FRONT_Y_SIGN > CHEST_NORMAL_Y
        )
        for loop in face.loops:
            if on_chest:
                co = loop.vert.co
                # Empirically calibrated for the -Y front so the wordmark reads
                # upright and left-to-right from the camera (see report: a plain
                # projection came out 180deg-rotated on this rig).
                u = (co.x - x0) / (x1 - x0)
                v = (co.z - z0) / (z1 - z0)
                loop[uvl].uv = (min(max(u, 0.0), 1.0), min(max(v, 0.0), 1.0))
            else:
                loop[uvl].uv = (2.0, 2.0)  # parked outside box
        if on_chest:
            placed += 1
    bm.to_mesh(me)
    bm.free()
    me.update()
    log(f"logo UV2 authored on {placed} chest faces")
    if placed == 0:
        log("WARNING: no chest faces matched — check CHEST_* box / front axis")


# --------------------------------------------------------------------------
# Export
# --------------------------------------------------------------------------

def export_reference(shell, armature, out_path):
    bpy.ops.object.select_all(action='DESELECT')
    shell.select_set(True)
    armature.select_set(True)
    bpy.context.view_layer.objects.active = armature
    bpy.ops.export_scene.gltf(
        filepath=out_path,
        export_format='GLB',
        use_selection=True,
        export_apply=False,       # keep Armature modifier for skinning
        export_animations=False,
        export_skins=True,
        export_yup=True,
        export_texcoords=True,
        export_normals=True,
        export_materials='EXPORT',
        export_image_format='AUTO',
    )
    size_kb = os.path.getsize(out_path) // 1024
    log(f"exported reference -> {out_path} ({size_kb} KB)")


def save_albedo_sidecar(albedo_img, out_path):
    albedo_img.filepath_raw = out_path
    albedo_img.file_format = 'PNG'
    albedo_img.save()
    log(f"saved base albedo -> {out_path}")


# --------------------------------------------------------------------------
# Entry
# --------------------------------------------------------------------------

def main():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    if len(argv) < 2:
        print("Usage: -- <bodies_dir>/average_m <out_dir> "
              "[--offset M] [--sleeve-frac F]")
        sys.exit(1)
    body_dir = argv[0]
    out_dir = argv[1]

    global OFFSET_M, SLEEVE_FRAC
    if "--offset" in argv:
        OFFSET_M = float(argv[argv.index("--offset") + 1])
    if "--sleeve-frac" in argv:
        SLEEVE_FRAC = float(argv[argv.index("--sleeve-frac") + 1])

    os.makedirs(out_dir, exist_ok=True)
    clear_scene()

    shell, armature = build_covered_mesh(body_dir)
    sleeve_cut(shell, armature)
    offset_outward(shell, OFFSET_M)
    solidify(shell, CLOTH_THICKNESS_M)

    albedo_img = make_base_albedo_image()
    assign_fabric_material(shell, albedo_img)

    author_logo_uv(shell)  # do UV2 before mask bake (mask uses UV0/active)
    bake_region_mask(shell, os.path.join(out_dir, "reference_mask.png"))
    save_albedo_sidecar(albedo_img, os.path.join(out_dir, "base_albedo.png"))

    export_reference(shell, armature, os.path.join(out_dir, "average_m.glb"))
    log("DONE")


if __name__ == "__main__":
    main()