import json
import pathlib
import sys

import bpy


CONSTRAINT_TYPE_CODES = {
    "IK": 3,
    "COPY_ROTATION": 8,
}


def flatten_matrix(matrix):
    return [value for row in matrix for value in row]


def flatten_matrix_column_major(matrix):
    return [matrix[row][column] for column in range(4) for row in range(4)]


def main(blend_path: str, output_path: str) -> None:
    bpy.ops.wm.open_mainfile(filepath=str(pathlib.Path(blend_path).resolve()), load_ui=False)
    scene = bpy.context.scene
    armature_object = bpy.data.objects.get("PoseConstraintArmatureObject")
    mesh_object = bpy.data.objects.get("PoseConstraintMeshObject")
    if armature_object is None or mesh_object is None:
        raise RuntimeError("pose constraint fixture objects are missing")

    frames = [1, 5, 10, 15]
    samples = []
    constraint_metadata = {}
    for pose_bone in armature_object.pose.bones:
        constraint_metadata[pose_bone.name] = [
            {
                "name": constraint.name,
                "typeCode": CONSTRAINT_TYPE_CODES.get(constraint.type),
                "influence": constraint.influence,
                "targetObject": constraint.target.name if constraint.target else None,
                "poleTargetObject": constraint.pole_target.name if constraint.type == "IK" and constraint.pole_target else None,
            }
            for constraint in pose_bone.constraints
        ]

    for frame in frames:
        scene.frame_set(frame)
        depsgraph = bpy.context.evaluated_depsgraph_get()
        evaluated_armature = armature_object.evaluated_get(depsgraph)
        evaluated_mesh_object = mesh_object.evaluated_get(depsgraph)
        mesh = evaluated_mesh_object.to_mesh(preserve_all_data_layers=True, depsgraph=depsgraph)
        try:
            bones = {
                pose_bone.name: flatten_matrix(pose_bone.matrix)
                for pose_bone in evaluated_armature.pose.bones
            }
            samples.append({
                "frame": frame,
                "bones": bones,
                "positions": [coordinate for vertex in mesh.vertices for coordinate in vertex.co],
            })
        finally:
            evaluated_mesh_object.to_mesh_clear()

    result = {
        "schemaVersion": 1,
        "fixture": pathlib.Path(blend_path).name,
        "evaluator": "Blender Depsgraph",
        "blenderVersion": bpy.app.version_string,
        "armature": armature_object.data.name,
        "armatureObject": armature_object.name,
        "mesh": mesh_object.data.name,
        "meshObject": mesh_object.name,
        "constraintMetadata": constraint_metadata,
        "meshBindMatrix": flatten_matrix_column_major(mesh_object.matrix_local),
        "armatureWorldMatrix": flatten_matrix_column_major(armature_object.matrix_world),
        "parentWorldMatrix": flatten_matrix_column_major(armature_object.parent.matrix_world),
        "frames": frames,
        "samples": samples,
        "tolerance": {"maxPositionError": 1e-5, "rmsPositionError": 1e-6, "maxMatrixError": 1e-5},
    }
    pathlib.Path(output_path).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    if "--" not in sys.argv or len(sys.argv) <= sys.argv.index("--") + 2:
        raise SystemExit("usage: blender -b --python generate-pose-constraint-golden.py -- input.blend output.json")
    arguments = sys.argv[sys.argv.index("--") + 1:]
    main(arguments[0], arguments[1])