"""Builds Convict_Body_Augments.blend: one skeleton, the suit torso, and four stages of arms, legs, torso kit and helmet.

Run headless:  blender -b --python build_convict.py [-- --no-fbx]
Writes the .blend, build/build_report.json (topology audit + clearance check) and, unless --no-fbx, the FBX files.
"""
import bpy, sys, os, json
from math import radians
from mathutils import Vector, Matrix

HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE)
from lib_geo import Builder, finish, part_trees, overlaps
import convict_parts as cp

OUT = os.path.dirname(HERE); BLEND = os.path.join(OUT, 'Convict_Body_Augments.blend')
ARGS = sys.argv[sys.argv.index('--') + 1:] if '--' in sys.argv else []

bpy.ops.wm.read_factory_settings(use_empty=True)
sc = bpy.context.scene; sc.unit_settings.system = 'METRIC'; sc.unit_settings.scale_length = 1.0


# ---------------------------------------------------------------- materials (flat PBR values, no textures)
def mk_mat(name, col, metal, rough, emit=None, coat=0.0, spec=0.5):
    m = bpy.data.materials.new('M_Convict_' + name); m.use_nodes = True
    b = next(n for n in m.node_tree.nodes if n.type == 'BSDF_PRINCIPLED')
    b.inputs['Base Color'].default_value = (*col, 1); b.inputs['Metallic'].default_value = metal; b.inputs['Roughness'].default_value = rough
    b.inputs['Specular IOR Level'].default_value = spec
    if emit: b.inputs['Emission Color'].default_value = (*emit[0], 1); b.inputs['Emission Strength'].default_value = emit[1]
    if coat: b.inputs['Coat Weight'].default_value = coat; b.inputs['Coat Roughness'].default_value = .06
    m.diffuse_color = (*col, 1); m.metallic = metal; m.roughness = rough
    return m


MAT = {
    cp.SUIT: mk_mat('Suit', (.014, .016, .019), 0.0, .78, spec=.2),   # insulated suit: the human part
    cp.DARK: mk_mat('AlloyGloss', (.012, .013, .016), .90, .26),      # gloss-black alloy shells (stages 1 and 4)
    cp.MATTE: mk_mat('AlloyMatte', (.016, .017, .019), 0.0, .52, spec=.3),   # matte-black coated shells (stage 2)
    cp.STEEL: mk_mat('Steel', (.11, .115, .125), 1.0, .36),           # joints, pins, discs
    cp.RED: mk_mat('PlateRed', (.38, .018, .014), .30, .38),          # painted plate
    cp.RUB: mk_mat('Rubber', (.010, .010, .011), 0.0, .90),           # hoses, pads, vents
    cp.VIS: mk_mat('Visor', (.004, .005, .007), 0.0, .07, coat=1.0),  # glass
    cp.GLOW: mk_mat('StatusLight', (.9, .03, .22), 0.0, .3, emit=((1.0, .03, .22), 2.5)),
    cp.CER: mk_mat('Ceramic', (.62, .60, .55), 0.0, .42),             # pale face plate and markings
    cp.GREY: mk_mat('AlloyGrey', (.045, .047, .052), .70, .46),       # gunmetal frame (stage 3) and plates (stage 4)
}

# ---------------------------------------------------------------- collections
root = bpy.data.collections.new('Convict'); sc.collection.children.link(root)


def coll(name, parent=root):
    c = bpy.data.collections.new(name); parent.children.link(c); return c


c_rig = coll('Rig'); c_core = coll('Body_Core')
c_stage = {st: coll(f'Stage{st}_{cp.STAGE_NAMES[st]}') for st in (1, 2, 3, 4)}

# ---------------------------------------------------------------- skeleton
J = cp.make_joints(); TABLE = cp.bone_table(J); BONES = [b[0] for b in TABLE]
ad = bpy.data.armatures.new('SK_Convict'); arm = bpy.data.objects.new('Armature', ad); c_rig.objects.link(arm)
ad.display_type = 'STICK'
bpy.context.view_layer.objects.active = arm; bpy.ops.object.mode_set(mode='EDIT')
for name, parent, head, tail in TABLE:
    eb = ad.edit_bones.new(name); eb.head = head; eb.tail = tail
    if parent: eb.parent = ad.edit_bones[parent]
bpy.ops.object.mode_set(mode='OBJECT')

# ---------------------------------------------------------------- modules
AUD = {}; tp = cp.torso_prof(); OBJ = {}


def build(name, fn, c):
    B = Builder(BONES, cp.MATS); fn(B)
    ob, a = finish(B, name, arm, c, MAT); AUD[name] = a; OBJ[name] = ob
    return ob


core_ob = build('SKM_Convict_Core', lambda B: cp.core(B, J, tp), c_core)
STAGE_OBS = {}
for st in (1, 2, 3, 4):
    c = c_stage[st]; obs = []
    obs.append(build(f'SKM_Convict_S{st}_Head', lambda B, st=st: cp.head(B, st, J), c))
    obs.append(build(f'SKM_Convict_S{st}_Torso', lambda B, st=st: cp.TORSO[st](B, J, tp), c))
    for sg, sn in ((1, 'L'), (-1, 'R')):
        def f_arm(B, st=st, sg=sg):
            with B.side(sg): cp.arm(B, st, J)

        def f_leg(B, st=st, sg=sg):
            with B.side(sg): cp.leg(B, st, J)
        obs.append(build(f'SKM_Convict_S{st}_Arm_{sn}', f_arm, c)); obs.append(build(f'SKM_Convict_S{st}_Leg_{sn}', f_leg, c))
    STAGE_OBS[st] = obs

# UV: every part was laid out in metres as it was lofted; pack each module into 0..1
packed = 0
for ob in OBJ.values():
    bpy.ops.object.select_all(action='DESELECT'); ob.select_set(True); bpy.context.view_layer.objects.active = ob
    try:
        bpy.ops.object.mode_set(mode='EDIT'); bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.uv.select_all(action='SELECT'); bpy.ops.uv.pack_islands(margin=0.002); packed += 1
    except Exception as e:
        print('UVPACK FAIL', ob.name, e)
    finally:
        bpy.ops.object.mode_set(mode='OBJECT')
print('UV packed', packed, 'of', len(OBJ))

# ---------------------------------------------------------------- relaxed display pose (frame 1); frame 0 = bind pose
bpy.context.view_layer.objects.active = arm


def rot(name, axis, deg):
    """Turn a bone about an axis given in bind-pose space (the hinge travels with the parent)."""
    pb = arm.pose.bones[name]; pb.rotation_mode = 'QUATERNION'
    rest = pb.bone.matrix_local.to_3x3()
    q = (rest.inverted() @ Matrix.Rotation(radians(deg), 3, Vector(axis).normalized()) @ rest).to_quaternion()
    pb.rotation_quaternion = q @ pb.rotation_quaternion


def both(base, axis, deg):
    ax = Vector(axis); rot(base + '_l', ax, deg); rot(base + '_r', (ax.x, -ax.y, -ax.z), deg)


for pb in arm.pose.bones:
    pb.rotation_mode = 'QUATERNION'; pb.keyframe_insert('rotation_quaternion', frame=0)
both('upperarm', (0, 1, 0), 23); both('upperarm', (1, 0, 0), -5); both('lowerarm', J['HX'], -16)
for nme, (frs, lens) in J['F'].items():
    for k in range(3):
        both(f'{nme}_0{k + 1}', Vector(frs[k].col[1][:3]), -6 if nme == 'thumb' else -(12, 18, 12)[k])
both('thigh', (0, 1, 0), -3.5); both('foot', (0, 1, 0), 3.5); both('foot', (0, 0, 1), 5); rot('head', (1, 0, 0), 4)
for pb in arm.pose.bones: pb.keyframe_insert('rotation_quaternion', frame=1)
if arm.animation_data and arm.animation_data.action: arm.animation_data.action.name = 'Convict_Pose_Relaxed'
sc.frame_start = 0; sc.frame_end = 1

# ---------------------------------------------------------------- clearance check: bind pose and relaxed pose
ALLOW = ('ball', 'socket', 'pin', 'hose', 'knuckle', 'cable', 'strap')
GEAR = ('_Core', '_Torso', '_Head')


def keep(r):
    n, oa, na, ba, ob_, nb, bb = r
    if 'torso_suit' in (na, nb):
        other = ob_ if na == 'torso_suit' else oa
        if other.endswith(GEAR): return False            # gear worn on the suit sits in it by design
    return True


REPORT = {}
for label, frame_no in (('bind', 0), ('relaxed', 1)):
    sc.frame_set(frame_no); dg = bpy.context.evaluated_depsgraph_get()
    for st in (1, 2, 3, 4):
        mods = {ob.name: part_trees(ob, dg) for ob in [core_ob] + STAGE_OBS[st]}
        res = []
        if label == 'bind':
            allp = [p for ps in mods.values() for p in ps]
            res = [r for r in overlaps(allp, allp, ALLOW) if keep(r)]
        else:
            names = list(mods)
            for i in range(len(names)):
                for k in range(i + 1, len(names)):
                    res += [r for r in overlaps(mods[names[i]], mods[names[k]], ALLOW) if keep(r)]
            res.sort(reverse=True)
        REPORT[f'{label}_S{st}'] = [dict(tri_pairs=r[0], a=f'{r[1]}:{r[2]}[{r[3]}]', b=f'{r[4]}:{r[5]}[{r[6]}]') for r in res[:40]]
        print(f'CLEARANCE {label} S{st}: {len(res)} crossing pairs')

tot = {st: sum(AUD[o.name]['tri_count'] for o in [core_ob] + STAGE_OBS[st]) for st in (1, 2, 3, 4)}
json.dump(dict(audit=AUD, stage_triangles=tot, bones=len(BONES), clearance=REPORT), open(os.path.join(HERE, 'build_report.json'), 'w'), indent=1)
for k, a in AUD.items():
    print(f"AUDIT {k}: v{a['verts']} tri{a['tri_count']} (q{a['quads']} t{a['tris']} n{a['ngons']}) open{a['open_edges']} zero{a['zero_faces']}")
print('STAGE TRIS', tot)

# ---------------------------------------------------------------- line-up: the four stages side by side as collection instances
c_line = coll('Lineup'); ASM = {}
for st in (1, 2, 3, 4):
    a = bpy.data.collections.new(f'ASM_Stage{st}'); a.children.link(c_core); a.children.link(c_stage[st]); ASM[st] = a
    e = bpy.data.objects.new(f'Lineup_Stage{st}_{cp.STAGE_NAMES[st]}', None); e.instance_type = 'COLLECTION'; e.instance_collection = a
    e.empty_display_size = .1; e.location = (1.3 + (st - 1) * 1.05, 0, 0); c_line.objects.link(e)

sc.frame_set(1)
vl = bpy.context.view_layer


def lc(c, node=None):
    node = node or vl.layer_collection
    if node.collection == c: return node
    for ch in node.children:
        r = lc(c, ch)
        if r: return r


for st in (2, 3, 4): lc(c_stage[st]).exclude = True          # stage 1 is worn at the origin; switch stages by toggling these
bpy.ops.wm.save_as_mainfile(filepath=BLEND)
print('SAVED', BLEND)

# ---------------------------------------------------------------- FBX (after the save; the scene is scaled for Unreal and not saved again)
if '--no-fbx' not in ARGS:
    for st in (2, 3, 4): lc(c_stage[st]).exclude = False
    lc(c_line).exclude = True
    sc.frame_set(0); ad.pose_position = 'REST'
    S = Matrix.Scale(100, 4); sc.unit_settings.scale_length = 0.01       # centimetres in the file, so Unreal reads scale 1 on every bone
    ad.transform(S)
    for ob in OBJ.values(): ob.data.transform(S)
    fdir = os.path.join(OUT, 'fbx'); os.makedirs(os.path.join(fdir, 'modules'), exist_ok=True)

    def export(path, obs):
        bpy.ops.object.select_all(action='DESELECT')
        for o in obs: o.select_set(True)
        bpy.context.view_layer.objects.active = arm
        bpy.ops.export_scene.fbx(filepath=path, use_selection=True, object_types={'ARMATURE', 'MESH'}, mesh_smooth_type='FACE', add_leaf_bones=False,
                                 bake_anim=False, use_armature_deform_only=True, primary_bone_axis='Y', secondary_bone_axis='X', axis_forward='-Z', axis_up='Y',
                                 apply_unit_scale=True, apply_scale_options='FBX_SCALE_NONE')
    for st in (1, 2, 3, 4):
        export(os.path.join(fdir, f'SK_Convict_Stage{st}_{cp.STAGE_NAMES[st]}.fbx'), [arm, core_ob] + STAGE_OBS[st])
        for ob in STAGE_OBS[st]: export(os.path.join(fdir, 'modules', ob.name + '.fbx'), [arm, ob])
    export(os.path.join(fdir, 'modules', core_ob.name + '.fbx'), [arm, core_ob])
    print('FBX written to', fdir)