"""
The maintenance robots of Anchorage, made ready once (2026-10-21). The user: "place couple robots 2 inside 2 outside 1 next
to one of the generators use this robot model asset in this path "D:\\SubnauticaCLONE\\Assets\\VaultCache\\Robotscob94d93332ad0V1"".
That folder is the Unreal pack Robot_scout_R_21, which lies in TheDescentUE/Content/Robot_scout_R_21 (never edited);
build\\ue_export_kit.ps1 -Root /Game/Robot_scout_R_21 -Do meshes,skeletal,materials,textures -Render
-Out models\\anchorage\\kit_export_robot wrote its skeletal mesh (UE4 mannequin skeleton, 1.79 m), its animations and its
textures out of the project.

Each robot here is that skeletal mesh posed and baked to a plain mesh: the layout preview has no use for the skeleton, and
Unreal gets the skeletal mesh, its materials and its animation from the layout (the collection's robot_* properties).
One collection per robot, "PROP | Robot scout R21 <name>", in source_fab/robot_scout/robot_scout.blend:

    ROBOTS   name: pose, colour scheme, what it is doing
    poses    idle   the pack's ThirdPersonIdle, first frame
             walk   ThirdPersonWalk at its longest stride
             reach  the idle pose with the right hand on a point in front of it (two-bone IK), the back and head bent to it
    schemes  the pack's four paint schemes: base (ochre), Color_02 (white), Color_03 (khaki), Color_04 (teal)

Textures: 2K copies in source_fab/robot_scout/textures (ImageMagick, see the README): colour and emissive as JPEG, the
normal maps with green flipped (Unreal's are DirectX, Blender reads OpenGL), the ORM maps (R occlusion, G roughness,
B metallic) as 8-bit PNG. Every robot faces -Y and stands on z = 0, like the other props.

    blender -b --factory-startup --python build/robot_prepare.py
"""
import math
import os
import sys

import bpy
from mathutils import Matrix, Vector

sys.dont_write_bytecode = True
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
KIT = os.path.join(ROOT, "kit_export_robot")
OUT = os.path.join(ROOT, "source_fab", "robot_scout")
TEX = os.path.join(OUT, "textures")
PACK = "/Game/Robot_scout_R_21"
UE_MESH = PACK + "/Mesh/SK_Robot_scout_R21"
UE_ANIM = {"idle": PACK + "/Demo/Animations/ThirdPersonIdle", "walk": PACK + "/Demo/Animations/ThirdPersonWalk"}
PARTS = ("hh", "legs", "Tors")          # the mesh's three material slots, in slot order
# scheme: (texture suffix, Unreal material folder, Unreal material names in slot order)
SCHEMES = {
    "ochre": ("BaseColor", "Base_color", ("M_Robo_hh_BaseColor_Mat", "M_robo_legs_BaseColor_Mat", "M_Robo_Tors_BaseColor_Mat")),
    "white": ("Color_02", "Color_02", ("M_Robo_hh_Color_02_Mat", "M_robo_legs_Color_02_Mat", "M_Robo_Tors_Color_02_Mat")),
    "khaki": ("Color_03", "Color_03", ("M_Robo_hh_Color_03_Mat", "M_robo_legs_Color_03_Mat", "M_Robo_Tors_Color_03_Mat")),
    # the pack's Color_04 folder names its head and legs materials ..._Color_03_Mat; they do use the Color_04 maps
    "teal": ("Color_04", "Color_04", ("M_Robo_hh_Color_03_Mat", "M_Robo_legs_Color_03_Mat", "M_Robo_Tors_Color_04_Mat")),
}
EMISSION = 4.0
# name: pose, scheme, the wrist's place for a reach (below), back and head bend (deg), what it is doing. Robot space:
# forward -Y, its right hand on the -X side, z up.
ROBOTS = {
    # wrist: where the wrist goes, from the right shoulder (after the back is bent): sideways, forward (-y), down (m).
    # The arm is 0.58 m from shoulder to wrist; about 0.49 m leaves the elbow a little bent.
    # The cooler's LOCK panel is 1.0-1.1 m above the round deck the robot stands on.
    "cooler": dict(pose="reach", scheme="ochre", wrist=(-0.05, -0.40, -0.28), bend=(8.0, 18.0),
                   note="outside: at work on an ice cooler's control screen"),
    "patrol": dict(pose="walk", scheme="khaki", note="outside: walking its round on the ice"),
    # the screen on top of the Scifi container 005 (1.20 m)
    "supply": dict(pose="reach", scheme="teal", wrist=(-0.05, -0.42, -0.24), bend=(14.0, 26.0),
                   note="inside: checking a supply container's screen"),
    "hall": dict(pose="walk", scheme="white", note="inside: walking through the station"),
}


def import_fbx(path):
    before = set(bpy.data.objects)
    bpy.ops.import_scene.fbx(filepath=path)
    return [o for o in bpy.data.objects if o not in before]


def take_action(path):
    """The action of an animation FBX; the armature it came on is thrown away."""
    new = import_fbx(path)
    act = None
    for o in new:
        if o.animation_data and o.animation_data.action:
            act = o.animation_data.action
    for o in new:
        bpy.data.objects.remove(o)
    # The action also keys the armature object (its 0.01 scale and placement from the file); the robot's armature has
    # its own, so only the bones' channels are kept. Of those, the bones' offsets are the UE4 mannequin's, for whom the
    # animations were made: applied to the robot they pull its toe plates off the feet. As Unreal's usual retargeting
    # does (translation from the skeleton, the pelvis's from the animation), only the rotations are kept, and the pelvis's
    # offset.
    for lay in act.layers:
        for st in lay.strips:
            for cb in st.channelbags:
                for fc in [f for f in cb.fcurves if not f.data_path.startswith("pose.bones")
                           or (f.data_path.endswith(".location") and '"pelvis"' not in f.data_path)]:
                    cb.fcurves.remove(fc)
    act.use_fake_user = True
    return act


def set_action(arm, act, frame):
    ad = arm.animation_data or arm.animation_data_create()
    ad.action = act
    if getattr(act, "slots", None) and len(act.slots):
        ad.action_slot = act.slots[0]
    bpy.context.scene.frame_set(int(frame), subframe=frame - int(frame))
    # keep the pose and let go of the action, so that hand-set bones are not overwritten
    basis = {pb.name: pb.matrix_basis.copy() for pb in arm.pose.bones}
    ad.action = None
    for pb in arm.pose.bones:
        pb.matrix_basis = basis[pb.name]
    bpy.context.view_layer.update()


def rot_world(arm, bone, axis, deg):
    """Turn a pose bone about a world axis through its head."""
    pb = arm.pose.bones[bone]
    ax = (arm.matrix_world.to_3x3().inverted() @ Vector(axis)).normalized()
    h = pb.head.copy()
    pb.matrix = Matrix.Translation(h) @ Matrix.Rotation(math.radians(deg), 4, ax) @ Matrix.Translation(-h) @ pb.matrix
    bpy.context.view_layer.update()


def turn_onto(arm, bone, v_from, v_to):
    """Turn a pose bone about its head by the shortest rotation that takes world direction v_from onto v_to."""
    pb = arm.pose.bones[bone]
    m3 = arm.matrix_world.to_3x3()
    r = m3.inverted() @ v_from.normalized().rotation_difference(v_to.normalized()).to_matrix() @ m3
    h = pb.head.copy()
    pb.matrix = Matrix.Translation(h) @ r.to_4x4() @ Matrix.Translation(-h) @ pb.matrix
    bpy.context.view_layer.update()


def reach(arm, target, pole):
    """Two-bone IK by hand for the right arm: the wrist (hand_r's head) onto `target`, the elbow bent towards `pole`
    (world points). Blender's IK constraint cannot be used: the FBX's bone tails do not lie on the next joint."""
    mw = arm.matrix_world
    pos = lambda b: mw @ arm.pose.bones[b].head          # noqa: E731
    s, e, w = pos("upperarm_r"), pos("lowerarm_r"), pos("hand_r")
    a, b = (e - s).length, (w - e).length
    t = Vector(target)
    d = min((t - s).length, (a + b) * 0.995)
    u = (t - s).normalized()
    p = Vector(pole) - s
    p = (p - u * p.dot(u)).normalized()
    cos_s = max(-1.0, min(1.0, (a * a + d * d - b * b) / (2.0 * a * d)))
    e2 = s + u * (a * cos_s) + p * (a * math.sqrt(1.0 - cos_s * cos_s))
    turn_onto(arm, "upperarm_r", e - s, e2 - s)
    e, w = pos("lowerarm_r"), pos("hand_r")
    turn_onto(arm, "lowerarm_r", w - e, (s + u * d) - e)
    return pos("hand_r"), d < (t - s).length - 1e-4


def longest_stride(arm, act):
    ad = arm.animation_data or arm.animation_data_create()
    ad.action = act
    if getattr(act, "slots", None) and len(act.slots):
        ad.action_slot = act.slots[0]
    best = (0.0, 1.0)
    f0, f1 = act.frame_range
    f = f0
    while f <= f1:
        bpy.context.scene.frame_set(int(f), subframe=f - int(f))
        a = arm.matrix_world @ arm.pose.bones["foot_l"].head
        b = arm.matrix_world @ arm.pose.bones["foot_r"].head
        d = abs(a.y - b.y)
        if d > best[0]:
            best = (d, f)
        f += 0.25
    ad.action = None
    return best[1], best[0]


def material(scheme, part):
    suffix = SCHEMES[scheme][0]
    name = "robot_scout M_R21_%s_%s" % (part, scheme)
    m = bpy.data.materials.get(name)
    if m:
        return m
    m = bpy.data.materials.new(name)
    m.use_nodes = True
    nt = m.node_tree
    nt.nodes.clear()
    out = nt.nodes.new("ShaderNodeOutputMaterial")
    out.location = (600, 0)
    bs = nt.nodes.new("ShaderNodeBsdfPrincipled")
    bs.location = (250, 0)
    nt.links.new(bs.outputs[0], out.inputs[0])

    def img(fname, x, y, colour):
        n = nt.nodes.new("ShaderNodeTexImage")
        n.image = bpy.data.images.load(os.path.join(TEX, fname), check_existing=True)
        n.image.colorspace_settings.name = "sRGB" if colour else "Non-Color"
        n.location = (x, y)
        return n
    col = img("T_R21_%s_%s.jpg" % (part, suffix), -700, 300, True)
    orm = img("T_R21_%s_ORM.png" % part, -700, 0, False)
    nrm = img("T_R21_%s_Normal_GL.png" % part, -700, -300, False)
    emi = img("T_R21_%s_Emissive.jpg" % part, -700, -600, True)
    sep = nt.nodes.new("ShaderNodeSeparateColor")
    sep.location = (-400, 0)
    nt.links.new(orm.outputs["Color"], sep.inputs[0])
    mul = nt.nodes.new("ShaderNodeMix")
    mul.data_type = "RGBA"
    mul.blend_type = "MULTIPLY"
    mul.location = (-150, 300)
    sock = {s.identifier: s for s in mul.inputs}          # the Mix node has an A and a B per data type
    sock["Factor_Float"].default_value = 1.0
    nt.links.new(col.outputs["Color"], sock["A_Color"])
    nt.links.new(sep.outputs["Red"], sock["B_Color"])
    nt.links.new(next(s for s in mul.outputs if s.identifier == "Result_Color"), bs.inputs["Base Color"])
    nt.links.new(sep.outputs["Green"], bs.inputs["Roughness"])
    nt.links.new(sep.outputs["Blue"], bs.inputs["Metallic"])
    nm = nt.nodes.new("ShaderNodeNormalMap")
    nm.location = (-150, -300)
    nt.links.new(nrm.outputs["Color"], nm.inputs["Color"])
    nt.links.new(nm.outputs["Normal"], bs.inputs["Normal"])
    nt.links.new(emi.outputs["Color"], bs.inputs["Emission Color"])
    bs.inputs["Emission Strength"].default_value = EMISSION
    return m


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    new = import_fbx(os.path.join(KIT, "meshes", "Mesh", "SK_Robot_scout_R21.fbx"))
    arm = next(o for o in new if o.type == "ARMATURE")
    body = next(o for o in new if o.type == "MESH")
    acts = {k: take_action(os.path.join(KIT, "anims", "Demo", os.path.basename(p) + ".fbx")) for k, p in UE_ANIM.items()}
    fps_src = float(bpy.context.scene.render.fps)          # the importer sets it from the files' time base
    stride_frame, stride = longest_stride(arm, acts["walk"])
    print("[robot] walk: longest stride %.2f m at frame %.2f" % (stride, stride_frame))
    target = bpy.data.objects.new("reach target", None)
    bpy.context.scene.collection.objects.link(target)
    keep = []
    for name, cfg in ROBOTS.items():
        for pb in arm.pose.bones:
            for c in list(pb.constraints):
                pb.constraints.remove(c)
            pb.matrix_basis = Matrix.Identity(4)
        if cfg["pose"] == "walk":
            set_action(arm, acts["walk"], stride_frame)
            anim, t = UE_ANIM["walk"], (stride_frame - 1.0) / fps_src
        else:
            set_action(arm, acts["idle"], 1.0)
            anim, t = UE_ANIM["idle"], 0.0
        if cfg["pose"] == "reach":
            back, head = cfg["bend"]
            for b, share in (("spine_01", 0.3), ("spine_02", 0.35), ("spine_03", 0.35)):
                rot_world(arm, b, (1, 0, 0), back * share)
            rot_world(arm, "neck_01", (1, 0, 0), head * 0.5)
            rot_world(arm, "head", (1, 0, 0), head * 0.5)
            sh = arm.matrix_world @ arm.pose.bones["upperarm_r"].head
            want = sh + Vector(cfg["wrist"])
            # the elbow out to the right side and down, as when working at a panel
            hand, short = reach(arm, want, sh + Vector((-0.45, 0.15, -0.55)))
            print("[robot] %s: wrist at (%.2f, %.2f, %.2f), wanted (%.2f, %.2f, %.2f)%s" % (
                name, hand.x, hand.y, hand.z, want.x, want.y, want.z, " (out of reach: arm straight)" if short else ""))
        dg = bpy.context.evaluated_depsgraph_get()
        ev = body.evaluated_get(dg)
        me = bpy.data.meshes.new_from_object(ev, preserve_all_data_layers=False, depsgraph=dg)
        me.transform(body.matrix_world)
        zmin = min(v.co.z for v in me.vertices)
        if abs(zmin) > 0.01:
            me.transform(Matrix.Translation((0, 0, -zmin)))
        me.name = "robot_%s" % name
        for i, part in enumerate(PARTS):
            me.materials[i] = material(cfg["scheme"], part)
        ob = bpy.data.objects.new("robot_%s Robot scout R21" % name, me)
        col = bpy.data.collections.new("PROP | Robot scout R21 %s" % name)
        bpy.context.scene.collection.children.link(col)
        col.objects.link(ob)
        folder, mats = SCHEMES[cfg["scheme"]][1], SCHEMES[cfg["scheme"]][2]
        col["prop_key"] = "robot_%s" % name
        col["prop_title"] = "Robot scout R21 %s" % name
        col["prop_source"] = "TheDescentUE/Content/Robot_scout_R_21 (the user's VaultCache/Robotscob94d93332ad0V1), posed and baked"
        col["prop_scale"] = 1.0
        col["robot_skeletal_mesh"] = UE_MESH
        col["robot_materials"] = ["%s/mats/Face_01/%s/%s" % (PACK, folder, m) for m in mats]
        col["robot_anim"] = anim
        col["robot_anim_time"] = round(t, 3)
        col["robot_pose"] = cfg["pose"]
        col["robot_scheme"] = cfg["scheme"]
        col["robot_note"] = cfg["note"]
        if cfg["pose"] == "reach":
            # the fingertips: the mesh's most forward point; placing the robot puts it on the panel it works at
            tip = min(me.vertices, key=lambda v: v.co.y).co
            col["robot_reach_point"] = [round(tip.x, 4), round(tip.y, 4), round(tip.z, 4)]
            col["robot_reach_wrist_from_shoulder"] = list(cfg["wrist"])
            col["robot_reach_bend"] = list(cfg["bend"])
            print("[robot] %s: fingertips at (%.2f, %.2f, %.2f)" % (name, tip.x, tip.y, tip.z))
        d = Vector([max(v.co[i] for v in me.vertices) - min(v.co[i] for v in me.vertices) for i in range(3)])
        print("[robot] %-7s %-5s %-6s %.2f x %.2f x %.2f m, %d triangles, anim %s at %.2f s" % (
            name, cfg["pose"], cfg["scheme"], d.x, d.y, d.z, sum(len(p.vertices) - 2 for p in me.polygons),
            anim.split("/")[-1], t))
        keep.append(ob)
    for o in list(bpy.data.objects):
        if o not in keep:
            bpy.data.objects.remove(o)
    for a in list(bpy.data.actions):
        bpy.data.actions.remove(a)
    for m in list(bpy.data.materials):
        if not m.users:
            bpy.data.materials.remove(m)
    for me in list(bpy.data.meshes):
        if not me.users:
            bpy.data.meshes.remove(me)
    for arm_data in list(bpy.data.armatures):
        bpy.data.armatures.remove(arm_data)
    path = os.path.join(OUT, "robot_scout.blend")
    bpy.ops.wm.save_as_mainfile(filepath=path, relative_remap=True, compress=True)
    print("[robot] saved %s" % path)


main()