# Robotic arm, stage 2: unwrap -> procedural wear/dirt materials -> bake to one texture set -> skeleton + skinning -> export.
# Run headless:  blender -b --factory-startup -P texture_and_rig.py
# Outputs next to this file: robotic-arm-rigged.blend / .glb / .fbx, textures/RA_*.png, preview/tex_*.png
#
# Rig notes: FK skeleton, UE-style names (root, upperarm_r, lowerarm_r, wrist_r, hand_r, <finger>_0N_r, piston_rod_r).
#   Bone Y runs along the part, bone X is the hinge: elbow bends with +X, hand flexes and fingers curl with -X,
#   wrist rolls on Y, the piston rod slides on Y. Each part stays its own mesh object, weighted 100% to its bone.
# Texture notes: 2048 px, one material. RA_ORM = R occlusion, G roughness, B metallic. RA_Normal is OpenGL (+Y);
#   in Unreal tick "Flip Green Channel" on that texture.
import bpy, bmesh, os, sys, time
from math import radians
from mathutils import Vector, Matrix
import numpy as np

D = os.path.dirname(os.path.abspath(__file__))
RES = 2048
t0 = time.time()
def log(*a): print('[RA]', round(time.time()-t0, 1), *a); sys.stdout.flush()

for o in list(bpy.data.objects): bpy.data.objects.remove(o)          # factory scene: drop cube/light/camera
exec(compile(open(os.path.join(D, 'build_robotic_arm.py'), encoding='utf-8').read(), 'build', 'exec'), {'__name__': '__main__'})
col = bpy.data.collections['RoboArm']; sc = bpy.context.scene; vl = bpy.context.view_layer
meshes = [o for o in col.objects if o.type == 'MESH']
log('built', len(meshes), 'meshes')

# ---------- triangulate (so baked tangents match the exported mesh) + unwrap into one atlas ----------
for o in meshes:
    bm = bmesh.new(); bm.from_mesh(o.data); bmesh.ops.triangulate(bm, faces=bm.faces[:]); bm.to_mesh(o.data); bm.free()
bpy.ops.object.select_all(action='DESELECT')
for o in meshes: o.select_set(True)
vl.objects.active = meshes[0]
bpy.ops.object.mode_set(mode='EDIT'); bpy.ops.mesh.select_all(action='SELECT')
bpy.ops.uv.smart_project(angle_limit=radians(66), island_margin=0.0025, area_weight=0.0, correct_aspect=True, scale_to_bounds=False)
bpy.ops.object.mode_set(mode='OBJECT')
G = 512; grid = np.zeros((G, G), np.int32); own = 0; shared = 0; oob = 0
for i, o in enumerate(meshes):                                         # crude cross-object overlap check
    uv = np.zeros(len(o.data.loops)*2, np.float32); o.data.uv_layers.active.data.foreach_get('uv', uv); uv = uv.reshape(-1, 2)
    oob += int(((uv < 0) | (uv > 1)).any(axis=1).sum())
    c = np.unique(np.clip((uv*G).astype(np.int32), 0, G-1), axis=0)
    prev = grid[c[:, 1], c[:, 0]]; shared += int((prev != 0).sum()); own += len(c); grid[c[:, 1], c[:, 0]] = i+1
log('uv cells', own, 'shared between objects', shared, 'loops outside 0-1', oob)

# ---------- procedural source materials (object-space, so no seams) ----------
SRC = {}
def build_src(m, base, metal, rough, wearcol=(0.13, 0.13, 0.14), wear=1.0, dirtcol=(0.10, 0.085, 0.07), dirt=0.75, ribs=0.0, scratch=0.2):
    nt = m.node_tree; nt.nodes.clear(); L = nt.links.new
    def node(t, **kw):
        n = nt.nodes.new(t)
        for k, v in kw.items(): setattr(n, k, v)
        return n
    def put(n, idx, x):
        if isinstance(x, (int, float)): n.inputs[idx].default_value = x
        elif isinstance(x, tuple): n.inputs[idx].default_value = (*x, 1)
        else: L(x, n.inputs[idx])
    def math(op, a, b=None, clamp=False):
        n = node('ShaderNodeMath', operation=op); n.use_clamp = clamp; put(n, 0, a)
        if b is not None: put(n, 1, b)
        return n.outputs[0]
    def mixc(f, a, b):
        n = node('ShaderNodeMix', data_type='RGBA'); n.clamp_factor = True; put(n, 0, f); put(n, 6, a); put(n, 7, b); return n.outputs[2]
    def mixf(f, a, b):
        n = node('ShaderNodeMix', data_type='FLOAT'); n.clamp_factor = True; put(n, 0, f); put(n, 2, a); put(n, 3, b); return n.outputs[0]
    def ramp(x, lo, hi):
        n = node('ShaderNodeMapRange', interpolation_type='SMOOTHSTEP'); put(n, 0, x); put(n, 1, lo); put(n, 2, hi); return n.outputs[0]
    def noise(vec, scale, detail=4.0, r=0.55):
        n = node('ShaderNodeTexNoise'); n.inputs['Scale'].default_value = scale; n.inputs['Detail'].default_value = detail
        n.inputs['Roughness'].default_value = r; L(vec, n.inputs['Vector']); return n.outputs[0]
    obj = node('ShaderNodeTexCoord').outputs['Object']
    mp = node('ShaderNodeMapping'); mp.inputs['Scale'].default_value = (5, 260, 260); L(obj, mp.inputs['Vector'])
    nbig, nwear, nfine, nchip = noise(obj, 14, 4), noise(obj, 70, 6, 0.65), noise(obj, 420, 3), noise(obj, 45, 5, 0.7)
    scr = ramp(noise(mp.outputs[0], 1.0, 2), 0.66, 0.8)                                   # long scratches along the part
    ao = node('ShaderNodeAmbientOcclusion', samples=16); ao.inputs['Distance'].default_value = 0.025; aoF = ao.outputs['AO']
    cav = math('SUBTRACT', 1.0, ramp(aoF, 0.35, 0.9))                                     # crevices
    bev = node('ShaderNodeBevel', samples=6); bev.inputs['Radius'].default_value = 0.0018
    dot = node('ShaderNodeVectorMath', operation='DOT_PRODUCT'); L(bev.outputs['Normal'], dot.inputs[0]); L(node('ShaderNodeNewGeometry').outputs['Normal'], dot.inputs[1])
    edge = math('MULTIPLY', math('SUBTRACT', 1.0, ramp(dot.outputs['Value'], 0.80, 0.995)), math('SUBTRACT', 1.0, cav))   # exposed edges only
    w_edge = ramp(math('MULTIPLY', edge, math('ADD', 0.45, nwear)), 0.55, 0.75)
    w = math('MULTIPLY', math('MAXIMUM', w_edge, math('MULTIPLY', ramp(nchip, 0.72, 0.78), 0.8)), wear)                 # chipped paint
    dmask = math('MULTIPLY', math('ADD', math('MULTIPLY', cav, math('ADD', 0.5, math('MULTIPLY', nbig, 1.2))), math('MULTIPLY', scr, scratch), clamp=True), dirt)
    c = mixc(nbig, tuple(v*0.86 for v in base), base)
    c = mixc(dmask, c, dirtcol); c = mixc(w, c, wearcol)
    r = math('ADD', math('ADD', rough, math('MULTIPLY', math('SUBTRACT', nfine, 0.5), 0.18)), math('MULTIPLY', dmask, 0.3), clamp=True)
    r = mixf(w, r, 0.42); mt = mixf(w, metal, 1.0)
    orm = node('ShaderNodeCombineColor'); L(aoF, orm.inputs[0]); L(r, orm.inputs[1]); L(mt, orm.inputs[2])
    h = math('SUBTRACT', math('SUBTRACT', math('MULTIPLY', nfine, 0.3), math('MULTIPLY', w, 0.6)), math('MULTIPLY', scr, 0.4*scratch/0.2))
    if ribs:
        wv = node('ShaderNodeTexWave', wave_type='BANDS', bands_direction='X'); wv.inputs['Scale'].default_value = 230; L(obj, wv.inputs['Vector'])
        h = math('ADD', h, math('MULTIPLY', wv.outputs[0], ribs))
    bump = node('ShaderNodeBump'); bump.inputs['Strength'].default_value = 0.6; bump.inputs['Distance'].default_value = 0.0006; L(h, bump.inputs['Height'])
    bsdf = node('ShaderNodeBsdfPrincipled'); L(c, bsdf.inputs['Base Color']); L(r, bsdf.inputs['Roughness']); L(mt, bsdf.inputs['Metallic']); L(bump.outputs[0], bsdf.inputs['Normal'])
    emit = node('ShaderNodeEmission'); out = node('ShaderNodeOutputMaterial'); img = node('ShaderNodeTexImage')
    SRC[m.name] = dict(nt=nt, col=c, orm=orm.outputs[0], bsdf=bsdf, emit=emit, out=out, img=img)

M_ = bpy.data.materials
build_src(M_['RA_ArmorWhite'], (0.80, 0.80, 0.78), 0.1, 0.35)
build_src(M_['RA_ArmorGrey'], (0.42, 0.43, 0.44), 0.2, 0.40)
build_src(M_['RA_Rubber'], (0.03, 0.04, 0.055), 0.0, 0.55, wear=0.0, dirtcol=(0.2, 0.19, 0.17), dirt=0.45, ribs=0.5, scratch=0.05)
build_src(M_['RA_Frame'], (0.07, 0.07, 0.075), 0.6, 0.45, wearcol=(0.55, 0.56, 0.58), dirtcol=(0.16, 0.13, 0.10), dirt=0.6)
build_src(M_['RA_Steel'], (0.55, 0.56, 0.58), 1.0, 0.25, wear=0.0, dirtcol=(0.08, 0.06, 0.045), dirt=0.85, scratch=0.5)
build_src(M_['RA_HoseRed'], (0.45, 0.03, 0.03), 0.0, 0.40, wear=0.0, dirtcol=(0.08, 0.03, 0.03), dirt=0.7, ribs=0.6, scratch=0.05)
build_src(M_['RA_Brass'], (0.6, 0.42, 0.15), 1.0, 0.30, wear=0.0, dirtcol=(0.10, 0.13, 0.09), dirt=0.8, scratch=0.4)
log('materials ready')

# ---------- bake ----------
sc.render.engine = 'CYCLES'
try:
    prefs = bpy.context.preferences.addons['cycles'].preferences
    for dt in ('OPTIX', 'CUDA', 'HIP', 'ONEAPI'):
        try:
            prefs.compute_device_type = dt; prefs.get_devices()
            devs = [d for d in prefs.devices if d.type == dt]
            if devs:
                for d in prefs.devices: d.use = d.type == dt
                sc.cycles.device = 'GPU'; log('GPU bake with', dt, devs[0].name); break
        except Exception: pass
except Exception as e: log('no GPU prefs', e)
os.makedirs(os.path.join(D, 'textures'), exist_ok=True)
IMG = {}
for key, name, data, fill in (('col', 'RA_BaseColor', False, (0.3, 0.3, 0.3, 1)), ('orm', 'RA_ORM', True, (1, 0.5, 0, 1)), ('nrm', 'RA_Normal', True, (0.5, 0.5, 1, 1))):
    im = bpy.data.images.new(name, RES, RES, alpha=False, is_data=data); im.generated_color = fill; IMG[key] = im
sc.render.bake.margin = 8; sc.render.bake.use_clear = False; sc.render.bake.normal_space = 'TANGENT'
bpy.ops.object.select_all(action='DESELECT')
for o in meshes: o.select_set(True)
vl.objects.active = meshes[0]
for key in ('col', 'orm', 'nrm'):
    for s in SRC.values():
        nt = s['nt']; s['img'].image = IMG[key]; nt.nodes.active = s['img']
        for l in list(s['out'].inputs[0].links) + list(s['emit'].inputs[0].links): nt.links.remove(l)
        if key == 'nrm': nt.links.new(s['bsdf'].outputs[0], s['out'].inputs[0])
        else: nt.links.new(s[key], s['emit'].inputs[0]); nt.links.new(s['emit'].outputs[0], s['out'].inputs[0])
    sc.cycles.samples = 4 if key == 'nrm' else 16
    bpy.ops.object.bake(type='NORMAL' if key == 'nrm' else 'EMIT', margin=8, use_clear=False)
    im = IMG[key]; im.filepath_raw = os.path.join(D, 'textures', im.name + '.png'); im.file_format = 'PNG'; im.save()
    log('baked', im.name)

# ---------- final single material ----------
fm = bpy.data.materials.new('RA_Arm')
if not fm.node_tree: fm.use_nodes = True
nt = fm.node_tree; nt.nodes.clear()
def tex(im):
    n = nt.nodes.new('ShaderNodeTexImage'); n.image = im; return n
bs = nt.nodes.new('ShaderNodeBsdfPrincipled'); outn = nt.nodes.new('ShaderNodeOutputMaterial'); nt.links.new(bs.outputs[0], outn.inputs[0])
nt.links.new(tex(IMG['col']).outputs['Color'], bs.inputs['Base Color'])
sep = nt.nodes.new('ShaderNodeSeparateColor'); nt.links.new(tex(IMG['orm']).outputs['Color'], sep.inputs[0])
nt.links.new(sep.outputs[1], bs.inputs['Roughness']); nt.links.new(sep.outputs[2], bs.inputs['Metallic'])
nm = nt.nodes.new('ShaderNodeNormalMap'); nt.links.new(tex(IMG['nrm']).outputs['Color'], nm.inputs['Color']); nt.links.new(nm.outputs[0], bs.inputs['Normal'])
for o in meshes:
    o.data.materials.clear(); o.data.materials.append(fm)
    for p in o.data.polygons: p.material_index = 0
for name in list(SRC): bpy.data.materials.remove(M_[name])

# ---------- skeleton ----------
vl.update()
BONES = [('RA_UpperArm', 'upperarm_r', 'root', 0.30), ('RA_Forearm', 'lowerarm_r', 'upperarm_r', 0.27), ('RA_PistonRod', 'piston_rod_r', 'lowerarm_r', 0.09),
         ('RA_Wrist', 'wrist_r', 'lowerarm_r', 0.022), ('RA_Hand', 'hand_r', 'wrist_r', 0.095)]
for f, Ls in (('Index', (0.042, 0.026, 0.023)), ('Middle', (0.046, 0.029, 0.024)), ('Ring', (0.043, 0.027, 0.023)), ('Pinky', (0.034, 0.021, 0.020)), ('Thumb', (0.042, 0.032, 0.028))):
    for i in (1, 2, 3): BONES.append((f'RA_{f}_{i}', f'{f.lower()}_0{i}_r', 'hand_r' if i == 1 else f'{f.lower()}_0{i-1}_r', Ls[i-1]))
MW = {o.name: o.matrix_world.copy() for o in meshes}
arm = bpy.data.objects.new('Armature', bpy.data.armatures.new('RA_Skeleton')); col.objects.link(arm)   # object named "Armature" so Unreal does not add it as an extra root bone
bpy.ops.object.select_all(action='DESELECT'); arm.select_set(True); vl.objects.active = arm
bpy.ops.object.mode_set(mode='EDIT')
eb = arm.data.edit_bones
b = eb.new('root'); b.head = (0, 0, 0); b.tail = (0.05, 0, 0); b.align_roll(Vector((0, 0, 1)))
for oname, bname, parent, ln in BONES:
    m = MW[oname]; r = m.to_3x3(); b = eb.new(bname); b.head = m.translation; b.tail = m.translation + r.col[0].normalized()*ln
    b.align_roll(r.col[2]); b.parent = eb[parent]; b.use_connect = False
bpy.ops.object.mode_set(mode='OBJECT')
root_empty = bpy.data.objects.get('RA_Root')
for oname, bname, parent, ln in BONES:                               # rigid skin: one group per part, weight 1
    o = bpy.data.objects[oname]; o.data.transform(MW[oname]); o.parent = arm; o.matrix_parent_inverse = Matrix.Identity(4); o.matrix_basis = Matrix.Identity(4)
    o.vertex_groups.new(name=bname).add(range(len(o.data.vertices)), 1.0, 'REPLACE')
    o.modifiers.new('Armature', 'ARMATURE').object = arm
if root_empty: bpy.data.objects.remove(root_empty)
arm.data.display_type = 'STICK'; arm.show_in_front = True
for pb in arm.pose.bones:
    pb.rotation_mode = 'XYZ'; n = pb.name
    if n == 'piston_rod_r': pb.lock_rotation = (True, True, True); pb.lock_location = (True, False, True)
    else:
        pb.lock_location = (True, True, True)
        if n == 'wrist_r': pb.lock_rotation = (True, False, True)
        elif n == 'lowerarm_r' or n[-4:] in ('02_r', '03_r'): pb.lock_rotation = (False, True, True)
        elif n == 'hand_r' or n.endswith('01_r'): pb.lock_rotation = (False, True, False)

# ---------- test animation: rest -> bent elbow + fist -> rest ----------
POSE = {'lowerarm_r': (55, 0, 0), 'wrist_r': (0, 30, 0), 'hand_r': (-15, 0, 0), 'thumb_02_r': (-40, 0, 0), 'thumb_03_r': (-35, 0, 0)}
for f in ('index', 'middle', 'ring', 'pinky'): POSE.update({f'{f}_01_r': (-60, 0, 0), f'{f}_02_r': (-75, 0, 0), f'{f}_03_r': (-50, 0, 0)})
pbs = arm.pose.bones
for frame, k in ((1, 0.0), (20, 1.0), (40, 0.0)):
    for n, rot in POSE.items():
        pbs[n].rotation_euler = [radians(a*k) for a in rot]; pbs[n].keyframe_insert('rotation_euler', frame=frame)
    pbs['piston_rod_r'].location = (0, -0.03*k, 0); pbs['piston_rod_r'].keyframe_insert('location', frame=frame)
if arm.animation_data and arm.animation_data.action: arm.animation_data.action.name = 'RA_Fist'
sc.frame_start = 1; sc.frame_end = 40
log('rigged', len(arm.data.bones), 'bones')

# ---------- previews ----------
sc.render.engine = 'BLENDER_EEVEE'
cam = bpy.data.objects['RA_Cam']
def look(ob, frm, to): ob.location = frm; ob.rotation_euler = (Vector(to)-Vector(frm)).to_track_quat('-Z', 'Y').to_euler()
os.makedirs(os.path.join(D, 'preview'), exist_ok=True)
try:
    for name, frame, frm, to, lens in (('tex_a', 1, (1.1, -0.7, 0.45), (0.38, 0, -0.02), 50), ('tex_b', 1, (0.95, 0.32, 0.22), (0.66, 0, -0.03), 50),
                                       ('tex_d', 1, (0.55, -0.42, 0.3), (0.42, 0, 0), 50), ('tex_pose', 20, (0.95, -0.75, 0.5), (0.38, 0, 0.1), 45)):
        sc.frame_set(frame); cam.data.lens = lens; look(cam, frm, to); sc.render.filepath = os.path.join(D, 'preview', name + '.png'); bpy.ops.render.render(write_still=True)
except Exception as e: log('preview render failed:', e)
sc.frame_set(1); cam.data.lens = 50; look(cam, (1.1, -0.7, 0.45), (0.38, 0, -0.02))

# ---------- save + export ----------
bpy.ops.wm.save_as_mainfile(filepath=os.path.join(D, 'robotic-arm-rigged.blend'))
bpy.ops.object.select_all(action='DESELECT'); arm.select_set(True)
for o in meshes: o.select_set(True)
vl.objects.active = arm
bpy.ops.export_scene.gltf(filepath=os.path.join(D, 'robotic-arm-rigged.glb'), export_format='GLB', use_selection=True, export_yup=True, export_cameras=False, export_lights=False, export_animations=True)
log('glb done')
bpy.ops.export_scene.fbx(filepath=os.path.join(D, 'robotic-arm-rigged.fbx'), use_selection=True, object_types={'ARMATURE', 'MESH'}, add_leaf_bones=False,
                         use_armature_deform_only=True, mesh_smooth_type='FACE', use_tspace=True, bake_anim=True, bake_anim_use_all_actions=False, bake_anim_use_nla_strips=False, path_mode='AUTO')
log('fbx done; tris', sum(len(o.data.polygons) for o in meshes))