"""Builds Convict_Body_Augments.blend (version 2): 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] [-- --stages 1]
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 []
STAGES = [int(c) for c in ARGS[ARGS.index('--stages') + 1]] if '--stages' in ARGS else [1, 2, 3, 4]
if '--out' in ARGS: BLEND = os.path.join(OUT, ARGS[ARGS.index('--out') + 1])

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
# Plain PBR values (these are what a game material needs) plus, for Blender renders only, worn edges, grime in the
# hollows, a rounded-edge highlight and uneven gloss. Those extras are procedural and are not baked to textures.
def mk_mat(name, col, metal, rough, emit=None, coat=0.0, spec=0.5, wear=0.0, wear_col=(.55, .56, .58), grime=0.5, fiber=False):
    m = bpy.data.materials.new('M_Convict_' + name); m.use_nodes = True; nt = m.node_tree
    b = next(n for n in nt.nodes if n.type == 'BSDF_PRINCIPLED'); N = nt.nodes.new; L = nt.links.new
    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
    if emit: return m
    tc = N('ShaderNodeTexCoord'); ns = N('ShaderNodeTexNoise'); ns.inputs['Scale'].default_value = 28; ns.inputs['Detail'].default_value = 6; L(tc.outputs['Object'], ns.inputs['Vector'])
    rr = N('ShaderNodeMapRange'); rr.inputs['From Min'].default_value = .3; rr.inputs['From Max'].default_value = .7
    rr.inputs['To Min'].default_value = max(rough - .07, .02); rr.inputs['To Max'].default_value = min(rough + .10, 1); L(ns.outputs['Fac'], rr.inputs['Value'])
    ao = N('ShaderNodeAmbientOcclusion'); ao.inputs['Distance'].default_value = .03; ao.samples = 8
    base = N('ShaderNodeMix'); base.data_type = 'RGBA'; base.inputs['A'].default_value = (col[0] * grime, col[1] * grime, col[2] * grime, 1); base.inputs['B'].default_value = (*col, 1)
    L(ao.outputs['AO'], base.inputs['Factor']); last_col = base.outputs['Result']; last_r = rr.outputs['Result']
    if wear > 0:
        geo = N('ShaderNodeNewGeometry'); ramp = N('ShaderNodeValToRGB'); ramp.color_ramp.elements[0].position = .52; ramp.color_ramp.elements[1].position = .60
        L(geo.outputs['Pointiness'], ramp.inputs['Fac'])
        sp = N('ShaderNodeTexNoise'); sp.inputs['Scale'].default_value = 140; sp.inputs['Detail'].default_value = 4; L(tc.outputs['Object'], sp.inputs['Vector'])
        mul = N('ShaderNodeMath'); mul.operation = 'MULTIPLY'; L(ramp.outputs['Color'], mul.inputs[0]); L(sp.outputs['Fac'], mul.inputs[1])
        amt = N('ShaderNodeMath'); amt.operation = 'MULTIPLY'; amt.inputs[1].default_value = wear * 2.0; amt.use_clamp = True; L(mul.outputs[0], amt.inputs[0])
        wc = N('ShaderNodeMix'); wc.data_type = 'RGBA'; wc.inputs['B'].default_value = (*wear_col, 1); L(amt.outputs[0], wc.inputs['Factor']); L(last_col, wc.inputs['A']); last_col = wc.outputs['Result']
        wm = N('ShaderNodeMix'); wm.data_type = 'FLOAT'; wm.inputs['A'].default_value = metal; wm.inputs['B'].default_value = 1.0; L(amt.outputs[0], wm.inputs['Factor']); L(wm.outputs['Result'], b.inputs['Metallic'])
    L(last_col, b.inputs['Base Color']); L(last_r, b.inputs['Roughness'])
    bev = N('ShaderNodeBevel'); bev.inputs['Radius'].default_value = .0018; bev.samples = 4; nrm = bev.outputs['Normal']
    if fiber:                                              # synthetic muscle: fine strands running along the limb
        wv = N('ShaderNodeTexWave'); wv.wave_type = 'BANDS'; wv.bands_direction = 'X'; wv.inputs['Scale'].default_value = 240; wv.inputs['Distortion'].default_value = 1.2
        wv.inputs['Detail'].default_value = 1; L(tc.outputs['UV'], wv.inputs['Vector'])
        bp = N('ShaderNodeBump'); bp.inputs['Strength'].default_value = .5; bp.inputs['Distance'].default_value = .002; L(wv.outputs['Fac'], bp.inputs['Height']); L(bev.outputs['Normal'], bp.inputs['Normal'])
        nrm = bp.outputs['Normal']
    L(nrm, b.inputs['Normal'])
    return m


MAT = {
    cp.SUIT: mk_mat('Suit', (.011, .012, .014), 0.0, .40, spec=.5, grime=.6),                        # insulated suit: the human part (black rubber)
    cp.FIBER: mk_mat('SyntheticMuscle', (.013, .014, .016), 0.0, .42, spec=.5, fiber=True),          # the limbs under their plates
    cp.DARK: mk_mat('AlloyGloss', (.010, .011, .014), .85, .22, wear=.5, wear_col=(.42, .44, .47)),   # gloss-black shells (stages 1 and 4)
    cp.MATTE: mk_mat('AlloyMatte', (.018, .019, .021), 0.0, .50, spec=.35, wear=.9),                  # matte-black coated shells (stage 2)
    cp.STEEL: mk_mat('Steel', (.10, .105, .115), 1.0, .34, wear=.4),                                  # joints, pins, discs
    cp.CHROME: mk_mat('BrightMetal', (.62, .64, .66), 1.0, .22),                                      # rings, buttons, tendons
    cp.RED: mk_mat('PlateRed', (.36, .016, .012), .25, .36, wear=1.0, wear_col=(.50, .50, .52)),      # painted plate
    cp.RUB: mk_mat('Rubber', (.010, .010, .011), 0.0, .88, spec=.3),                                  # hoses, pads, vents
    cp.VIS: mk_mat('Visor', (.004, .005, .007), 0.0, .06, coat=1.0),                                  # glass
    cp.GLOW: mk_mat('StatusLight', (.9, .03, .22), 0.0, .3, emit=((1.0, .03, .22), 3.5)),
    cp.CER: mk_mat('Ceramic', (.60, .58, .53), 0.0, .40, wear=.5, wear_col=(.2, .2, .2)),             # pale face plate and markings
    cp.GREY: mk_mat('AlloyGrey', (.040, .042, .047), .75, .44, wear=.8),                              # 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 STAGES}

# ---------------------------------------------------------------- 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 = {}; 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), c_core)
STAGE_OBS = {}
for st in STAGES:
    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_kit(B, st, J), 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
    print('BUILT stage', st)

# UV: every part was laid out in metres as it was made; 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))

# ---------------------------------------------------------------- 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, dl=0.0):
    ax = Vector(axis); rot(base + '_l', ax, deg + dl); rot(base + '_r', (ax.x, -ax.y, -ax.z), deg - dl)


for pb in arm.pose.bones:
    pb.rotation_mode = 'QUATERNION'; pb.keyframe_insert('rotation_quaternion', frame=0)
both('upperarm', (0, 1, 0), 25); both('upperarm', (1, 0, 0), -7, 3); both('lowerarm', J['HX'], -20, 4)
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('thigh_l', (1, 0, 0), -5); rot('calf_l', (1, 0, 0), 8); rot('foot_l', (1, 0, 0), -3)            # weight on the right leg
rot('head', (1, 0, 0), 4); rot('head', (0, 0, 1), 7)
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 on the cages: bind pose and display pose
ALLOW = ('ball', 'socket', 'pin', 'hose', 'knuckle', 'cable', 'strap')
for ob in OBJ.values(): ob.modifiers['Subdivision'].show_viewport = False
REPORT = {}
for label, frame_no in (('bind', 0), ('relaxed', 1)):
    sc.frame_set(frame_no); dg = bpy.context.evaluated_depsgraph_get()
    for st in STAGES:
        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 = overlaps(allp, allp, ALLOW)
        else:
            names = list(mods)
            for i in range(len(names)):
                for k in range(i + 1, len(names)): res += overlaps(mods[names[i]], mods[names[k]], ALLOW)
            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')
for ob in OBJ.values(): ob.modifiers['Subdivision'].show_viewport = True

sc.frame_set(0); dg = bpy.context.evaluated_depsgraph_get(); GAME = {}
for ob in OBJ.values():
    me = ob.evaluated_get(dg).to_mesh(); GAME[ob.name] = sum(len(p.vertices) - 2 for p in me.polygons); ob.evaluated_get(dg).to_mesh_clear()
tot = {st: GAME[core_ob.name] + sum(GAME[o.name] for o in STAGE_OBS[st]) for st in STAGES}
cage = {st: AUD[core_ob.name]['tri_count'] + sum(AUD[o.name]['tri_count'] for o in STAGE_OBS[st]) for st in STAGES}
json.dump(dict(audit=AUD, game_triangles_per_module=GAME, stage_triangles_game=tot, stage_triangles_cage=cage, bones=len(BONES), clearance=REPORT),
          open(os.path.join(HERE, 'build_report.json'), 'w'), indent=1)
bad = {k: (a['ngons'], a['bad_edges'], a['zero_faces']) for k, a in AUD.items() if a['ngons'] or a['bad_edges'] or a['zero_faces']}
print('AUDIT modules', len(AUD), 'with n-gons / non-manifold edges / zero faces:', bad)
print('STAGE TRIS cage', cage, 'game (one subdivision)', tot)

# ---------------------------------------------------------------- line-up: the stages side by side as collection instances
c_line = coll('Lineup')
for i, st in enumerate(STAGES):
    a = bpy.data.collections.new(f'ASM_Stage{st}'); a.children.link(c_core); a.children.link(c_stage[st])
    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.4 + i * 1.15, 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 STAGES[1:]: lc(c_stage[st]).exclude = True         # the first stage 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 STAGES[1:]: 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)
    def export(path, obs, subdivided):
        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', use_mesh_modifiers=subdivided, use_mesh_modifiers_render=False)
    # two sets: 'fbx' is subdivided once (smooth, for close views such as the first-person arms), 'fbx_low' is the cage as modelled
    for sub, folder in ((True, 'fbx'), (False, 'fbx_low')):
        fdir = os.path.join(OUT, folder); os.makedirs(os.path.join(fdir, 'modules'), exist_ok=True)
        for st in STAGES:
            export(os.path.join(fdir, f'SK_Convict_Stage{st}_{cp.STAGE_NAMES[st]}.fbx'), [arm, core_ob] + STAGE_OBS[st], sub)
            for ob in STAGE_OBS[st]: export(os.path.join(fdir, 'modules', ob.name + '.fbx'), [arm, ob], sub)
        export(os.path.join(fdir, 'modules', core_ob.name + '.fbx'), [arm, core_ob], sub)
        print('FBX written to', fdir)