bl_info = {
    "name": "Avatar Workbench",
    "author": "Machinimatrix",
    "version": (1, 0, 0),
    "blender": (3, 6, 0),
    "location": "View3D > Sidebar > sl > Workbench",
    "description": "Add standalone Second Life reference models and tools",
    "category": "3D View",
}

import os
import xml.etree.ElementTree as ET
from struct import unpack

import bpy
from bpy.props import BoolProperty, PointerProperty
from mathutils import Vector


ADDON_DIR = os.path.dirname(__file__)
DATA_DIR = os.path.join(ADDON_DIR, "lib")
SKELETON_FILE = os.path.join(DATA_DIR, "avatar_skeleton.xml")
LAD_FILE = os.path.join(DATA_DIR, "avatar_lad.xml")


WEIGHTSMAP = {
    "hairMesh": [("mNeck", "mHead"), ("mHead", None)],
    "headMesh": [("mNeck", "mHead"), ("mHead", None)],
    "eyelashMesh": [("mHead", None)],
    "upperBodyMesh": [
        ("mPelvis", "mTorso"), ("mTorso", "mChest"), ("mChest", "mNeck"), ("mNeck", None),
        ("mChest", "mCollarLeft"), ("mCollarLeft", "mShoulderLeft"), ("mShoulderLeft", "mElbowLeft"),
        ("mElbowLeft", "mWristLeft"), ("mWristLeft", None), ("mChest", "mCollarRight"),
        ("mCollarRight", "mShoulderRight"), ("mShoulderRight", "mElbowRight"),
        ("mElbowRight", "mWristRight"), ("mWristRight", None),
    ],
    "lowerBodyMesh": [
        ("mPelvis", "mHipRight"), ("mHipRight", "mKneeRight"), ("mKneeRight", "mAnkleRight"),
        ("mAnkleRight", None), ("mPelvis", "mHipLeft"), ("mHipLeft", "mKneeLeft"),
        ("mKneeLeft", "mAnkleLeft"), ("mAnkleLeft", None),
    ],
    "skirtMesh": [
        ("mTorso", "mPelvis"), ("mPelvis", None), ("mPelvis", "mHipRight"),
        ("mHipRight", "mKneeRight"), ("mKneeRight", None), ("mPelvis", "mHipLeft"),
        ("mHipLeft", "mKneeLeft"), ("mKneeLeft", None),
    ],
    "eyeBallLeftMesh": [("mEyeLeft", None)],
    "eyeBallRightMesh": [("mEyeRight", None)],
}


def s2b(value):
    return Vector((value[1], -value[0], value[2]))


def float_array(value):
    return [float(v) for v in value.split()] if value else [0.0, 0.0, 0.0]


def bool_attr(element, name, default=False):
    value = element.get(name)
    if value is None:
        return default
    return value.lower() == "true"


class BoneDef:
    def __init__(self, name, parent_name, head, tail, connected=False, bone_type="DEFORM"):
        self.name = name
        self.parent_name = parent_name
        self.head = head
        self.tail = tail
        self.connected = connected
        self.bone_type = bone_type


def iter_skeleton_bones(root):
    defs = []

    def visit(element, parent_name, parent_head):
        for child in element:
            if child.tag not in {"bone", "collision_volume"}:
                continue

            name = child.get("name")
            pos = s2b(Vector(float_array(child.get("pivot") or child.get("pos"))))
            end = s2b(Vector(float_array(child.get("end"))))
            head = parent_head + pos
            tail = head + end
            if (tail - head).length < 0.0001:
                tail = head + Vector((0.0, 0.02, 0.0))
            connected = bool_attr(child, "connected")
            bone_type = "VOLUME" if child.tag == "collision_volume" else "DEFORM"
            defs.append(BoneDef(name, parent_name, head, tail, connected, bone_type))
            visit(child, name, head)

    visit(root, None, Vector((0.0, 0.0, 0.0)))
    return defs


def load_attachments(existing_bones):
    ladxml = ET.parse(LAD_FILE)
    skel = ladxml.find("skeleton")
    attachments = []
    if skel is None:
        return attachments

    for attach in skel.findall("attachment_point"):
        joint = attach.get("joint")
        if joint not in existing_bones and joint != "mRoot":
            continue
        name = "a" + attach.get("name")
        pos = s2b(Vector(float_array(attach.get("position"))))
        parent_name = "mPelvis" if joint == "mRoot" else joint
        parent = existing_bones.get(parent_name)
        if parent is None:
            continue
        head = parent.head + pos
        tail = head + Vector((0.0, 0.0, 0.03))
        attachments.append(BoneDef(name, parent_name, head, tail, False, "ATTACHMENT"))
    return attachments


def create_armature(context, include_deform=True, include_volume=False, include_attachments=False):
    root = ET.parse(SKELETON_FILE).getroot()
    skeleton_defs = iter_skeleton_bones(root)
    bone_map = {bone.name: bone for bone in skeleton_defs}
    bone_defs = [
        bone for bone in skeleton_defs
        if (
            bone.bone_type == "DEFORM" and include_deform
            or bone.bone_type == "VOLUME" and include_volume
        )
    ]
    if include_attachments:
        attachments = load_attachments(bone_map)
        bone_defs.extend(attachments)

    arm_data = bpy.data.armatures.new("SL Reference Skeleton")
    arm_obj = bpy.data.objects.new("SL Reference Model", arm_data)
    context.collection.objects.link(arm_obj)
    context.view_layer.objects.active = arm_obj
    arm_obj.select_set(True)

    bpy.ops.object.mode_set(mode="EDIT")
    edit_bones = arm_data.edit_bones
    for bone_def in bone_defs:
        bone = edit_bones.new(bone_def.name)
        bone.head = bone_def.head
        bone.tail = bone_def.tail
        if bone_def.parent_name and bone_def.parent_name in edit_bones:
            bone.parent = edit_bones[bone_def.parent_name]
            bone.use_connect = bone_def.connected and (bone.head - bone.parent.tail).length < 0.0001
        bone["sl_reference_type"] = bone_def.bone_type

    bpy.ops.object.mode_set(mode="OBJECT")
    arm_data.display_type = "STICK"
    arm_obj.show_in_front = True
    return arm_obj


def get_vertex_index(mesh, vertex):
    mapped = mesh["vertexRemap"].get(vertex, vertex)
    return mesh["vertLookup"].index(mapped)


def load_llm(name, filepath):
    with open(filepath, "rb") as stream:
        mesh = {}
        mesh["header"] = stream.read(24).decode("utf-8").split("\x00")[0]
        has_weights = unpack("B", stream.read(1))[0]
        has_detail_tex_coords = unpack("B", stream.read(1))[0]
        mesh["position"] = unpack("<3f", stream.read(12))
        mesh["rotationAngles"] = unpack("<3f", stream.read(12))
        mesh["rotationOrder"] = unpack("B", stream.read(1))[0]
        mesh["scale"] = unpack("<3f", stream.read(12))
        num_vertices = unpack("<H", stream.read(2))[0]

        scale = (1.0, 1.0, 1.0)
        shift = (0.0, 0.0, 0.0)
        if name == "eyeBallLeftMesh":
            shift = (0.0736, 0.0360, 1.7617)
        elif name == "eyeBallRightMesh":
            shift = (0.0736, -0.0360, 1.7617)

        mesh["baseCoords"] = [
            s2b(Vector((co[0] * scale[0] + shift[0], co[1] * scale[1] + shift[1], co[2] * scale[2] + shift[2])))
            for co in (unpack("<3f", stream.read(12)) for _ in range(num_vertices))
        ]
        mesh["baseNormals"] = [s2b(Vector(unpack("<3f", stream.read(12)))) for _ in range(num_vertices)]
        mesh["baseBinormals"] = [s2b(Vector(unpack("<3f", stream.read(12)))) for _ in range(num_vertices)]
        mesh["texCoords"] = [unpack("<2f", stream.read(8)) for _ in range(num_vertices)]

        if has_detail_tex_coords:
            mesh["detailTexCoords"] = [unpack("<2f", stream.read(8)) for _ in range(num_vertices)]

        if has_weights:
            mesh["weights"] = []
            for _ in range(num_vertices):
                raw = unpack("<f", stream.read(4))[0]
                mesh["weights"].append((int(raw) - 1, raw - int(raw)))
        elif name in {"eyeBallLeftMesh", "eyeBallRightMesh"}:
            mesh["weights"] = [(0, 0.0)] * num_vertices

        num_faces = unpack("<H", stream.read(2))[0]
        mesh["faces"] = [unpack("<3H", stream.read(6)) for _ in range(num_faces)]

        if has_weights:
            num_skin_joints = unpack("<H", stream.read(2))[0]
            mesh["skinJoints"] = [
                stream.read(64).decode("utf-8").split("\x00")[0]
                for _ in range(num_skin_joints)
            ]
        elif name == "eyeBallLeftMesh":
            mesh["skinJoints"] = ["mEyeLeft"]
        elif name == "eyeBallRightMesh":
            mesh["skinJoints"] = ["mEyeRight"]

        morph_name = stream.read(64).decode("utf-8").split("\x00")[0]
        while morph_name != "End Morphs":
            num_morph_vertices = unpack("<L", stream.read(4))[0]
            stream.seek(num_morph_vertices * 48, os.SEEK_CUR)
            morph_name = stream.read(64).decode("utf-8").split("\x00")[0]

        num_remaps = unpack("<l", stream.read(4))[0]
        vertex_remap = {}
        for _ in range(num_remaps):
            original, mapped = unpack("<2l", stream.read(8))
            vertex_remap[original] = mapped
        mesh["vertexRemap"] = vertex_remap
        mesh["vertLookup"] = [i for i in range(len(mesh["baseCoords"])) if i not in vertex_remap]
        mesh["name"] = name
        return mesh


def load_meshes():
    meshes = {}
    ladxml = ET.parse(LAD_FILE)
    for mesh_elem in ladxml.findall("mesh"):
        if int(mesh_elem.get("lod", "0")) != 0:
            continue
        name = mesh_elem.get("type")
        filepath = os.path.join(DATA_DIR, mesh_elem.get("file_name"))
        meshes[name] = load_llm(name, filepath)
    return meshes


def create_mesh_object(context, arm_obj, mesh_data):
    verts = [mesh_data["baseCoords"][i] for i in mesh_data["vertLookup"]]
    faces = [[get_vertex_index(mesh_data, vi) for vi in face] for face in mesh_data["faces"]]
    mesh = bpy.data.meshes.new(mesh_data["name"])
    mesh.from_pydata(verts, [], faces)
    mesh.update(calc_edges=True)

    obj = bpy.data.objects.new(mesh_data["name"], mesh)
    context.collection.objects.link(obj)
    obj.parent = arm_obj

    uv_layer = mesh.uv_layers.new(name="SLMap")
    for polygon in mesh.polygons:
        for loop_index in polygon.loop_indices:
            vi = mesh_data["vertLookup"][mesh.loops[loop_index].vertex_index]
            uv = mesh_data["texCoords"][vi]
            uv_layer.data[loop_index].uv = (uv[0], 1.0 - uv[1])

    for poly in mesh.polygons:
        poly.use_smooth = True

    if "skinJoints" in mesh_data:
        for joint in mesh_data["skinJoints"]:
            obj.vertex_groups.new(name=joint)
        weight_map = WEIGHTSMAP.get(mesh_data["name"], [])
        for source_index, (bone_index, weight) in enumerate(mesh_data["weights"]):
            if source_index in mesh_data["vertexRemap"] or bone_index < 0 or bone_index >= len(weight_map):
                continue
            target_index = get_vertex_index(mesh_data, source_index)
            first, second = weight_map[bone_index]
            first_group = obj.vertex_groups.get(first)
            second_group = obj.vertex_groups.get(second) if second else None
            if first_group:
                first_group.add([target_index], 1.0 - weight, "REPLACE")
            if second_group and weight:
                second_group.add([target_index], weight, "REPLACE")

    modifier = obj.modifiers.new("Armature", "ARMATURE")
    modifier.object = arm_obj
    return obj


def create_reference_meshes(context, arm_obj):
    container = bpy.data.collections.new("SL Reference Meshes")
    context.scene.collection.children.link(container)
    for mesh_data in load_meshes().values():
        obj = create_mesh_object(context, arm_obj, mesh_data)
        for coll in list(obj.users_collection):
            coll.objects.unlink(obj)
        container.objects.link(obj)


class AvatarWorkbenchProperties(bpy.types.PropertyGroup):
    deform_bones: BoolProperty(
        name="Deform Bones",
        description="Create the SL deform skeleton",
        default=True,
    )
    attachment_bones: BoolProperty(
        name="Attachment Bones",
        description="Create attachment point bones from avatar_lad.xml",
        default=False,
    )
    volume_bones: BoolProperty(
        name="Volume Bones",
        description="Create collision volume bones from avatar_skeleton.xml",
        default=False,
    )


class AVATAR_WORKBENCH_OT_add_reference_model(bpy.types.Operator):
    bl_idname = "object.add_avatar_workbench_reference_model"
    bl_label = "SL Reference Model"
    bl_description = "Add a standalone Second Life reference skeleton and optional system meshes"
    bl_options = {"REGISTER", "UNDO"}

    with_meshes: BoolProperty(
        name="With Meshes",
        description="Create the SL system reference meshes",
        default=True,
    )

    def execute(self, context):
        props = context.scene.avatar_workbench
        if not (props.deform_bones or props.attachment_bones or props.volume_bones):
            self.report({"ERROR"}, "Select at least one bone type.")
            return {"CANCELLED"}

        arm_obj = create_armature(
            context,
            include_deform=props.deform_bones,
            include_volume=props.volume_bones,
            include_attachments=props.attachment_bones,
        )
        if self.with_meshes:
            create_reference_meshes(context, arm_obj)
        context.view_layer.objects.active = arm_obj
        arm_obj.select_set(True)
        return {"FINISHED"}


class AVATAR_WORKBENCH_PT_workbench(bpy.types.Panel):
    bl_label = "Workbench"
    bl_idname = "AVATAR_WORKBENCH_PT_workbench"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "sl"

    def draw(self, context):
        layout = self.layout
        props = context.scene.avatar_workbench

        col = layout.column(align=True)
        col.prop(props, "deform_bones")
        col.prop(props, "attachment_bones")
        col.prop(props, "volume_bones")

        if not (props.deform_bones or props.attachment_bones or props.volume_bones):
            layout.label(text="Select at least one bone type.", icon="ERROR")

        layout.operator(AVATAR_WORKBENCH_OT_add_reference_model.bl_idname, text="Add SL Reference Model")


classes = (
    AvatarWorkbenchProperties,
    AVATAR_WORKBENCH_OT_add_reference_model,
    AVATAR_WORKBENCH_PT_workbench,
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    bpy.types.Scene.avatar_workbench = PointerProperty(type=AvatarWorkbenchProperties)


def unregister():
    del bpy.types.Scene.avatar_workbench
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)


if __name__ == "__main__":
    register()
