ybnimport refactor 1

from traceback import format_exc
import bpy

from typing import Union

from ..tools.utils import get_distance_of_vectors
from .collision_materials import create_collision_material_from_index
from ..tools.meshhelper import create_disc, create_vertexcolor_layer, create_box, get_direction_of_vectors
from ..resources.bound import *
from ..sollumz_properties import *
from .properties import BoundFlags, CollisionMatFlags
from numpy import absolute
import os


def apply_bound_properties(bpy_object: bpy.types.Object, bound_xml: Bound):
    """Sets the bound_properties data block of bound_bpy to the corresponding properties of bound_xml."""
    bpy_object.bound_properties.procedural_id = bound_xml.procedural_id
    bpy_object.bound_properties.room_id = bound_xml.room_id
    bpy_object.bound_properties.ped_density = bound_xml.ped_density
    bpy_object.bound_properties.poly_flags = bound_xml.poly_flags
    bpy_object.bound_properties.inertia = bound_xml.inertia
    bpy_object.bound_properties.unk_flags = bound_xml.unk_flags
    bpy_object.bound_properties.margin = bound_xml.margin
    bpy_object.bound_properties.volume = bound_xml.volume


def apply_bound_child_properties(bpy_object: bpy.types.Object, bound_xml: Bound):
    """Applys transforms and flags from bounds that are children of composites to the specified bpy_object."""
    bpy_object.matrix_world = bound_xml.composite_transform.transposed()
    assign_composite_flags(bpy_object.composite_flags1,
                           bound_xml.composite_flags1)
    assign_composite_flags(bpy_object.composite_flags2,
                           bound_xml.composite_flags2)


# TODO: Test
def assign_composite_flags(bpy_flags: BoundFlags, xml_flags: list[str]):
    """Assigns the list of xml_flags to the provided BoundFlags data block."""
    for flag in xml_flags:
        flag_lower = flag.lower()
        if hasattr(bpy_flags, flag_lower):
            setattr(bpy_flags, flag_lower, True)


# TODO: Make import and operators create bound shapes the same way using the same function
def bound_to_bpy_object(bound_xml: Bound, name: str) -> bpy.types.Object:
    """Create a bpy mesh object from generic cwxml bound object."""
    # Create blender objects
    mesh = bpy.meshes.new(name)
    bpy_object = bpy.data.objects.new(name, mesh)

    # Create materials
    mat_index = bpy_object.material_index
    try:
        # TODO: Remove try except, just prevent error from happening in the first place
        # TODO: What about the material properties?
        mat = create_collision_material_from_index(mat_index)
        mesh.materials.append(mat)
    except IndexError:
        print(
            f"Warning: Failed to set materials for {name}. Material index {mat_index} does not exist in collisionmats list.")

    apply_bound_properties(bpy_object, bound_xml)

    bpy.context.collection.objects.link(bpy_object)

    return bpy_object


def bound_to_bpy_empty(bound_xml: Bound, name: str) -> bpy.types.Object:
    """Create a bpy empty from generic cwxml Bound object with bound properties applied."""
    bpy_empty = bpy.data.objects.new(name, None)
    bpy_empty.empty_display_size = 0
    apply_bound_properties(bpy_empty, bound_xml)

    bpy.context.collection.objects.link(bpy_empty)

    return bpy_empty


def bound_box_to_bpy(box_xml: BoundBox) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml BoundBox object."""
    sollum_type = SollumType.BOUND_BOX

    box_bpy = bound_to_bpy_object(
        box_xml, name=SOLLUMZ_UI_NAMES[sollum_type])
    box_bpy.sollum_type

    apply_bound_child_properties(box_bpy, box_xml)

    # TODO: Make sure absolute() works as expected
    box_bpy.bound_dimensions = absolute(
        box_xml.box_max - box_xml.box_min)
    box_bpy.data.transform(Matrix.Translation(box_xml.box_center))

    return box_bpy


def bound_sphere_to_bpy(sphere_xml: BoundSphere) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml BoundSphere object."""
    sollum_type = SollumType.BOUND_SPHERE

    sphere_bpy = bound_to_bpy_object(
        sphere_xml, name=SOLLUMZ_UI_NAMES[sollum_type])
    sphere_bpy.sollum_type = sollum_type

    apply_bound_child_properties(sphere_bpy, sphere_xml)

    sphere_bpy.bound_radius = sphere_xml.sphere_radius

    return sphere_bpy


def bound_capsule_to_bpy(capsule_xml: BoundCapsule) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml BoundCapsule object."""
    sollum_type = SollumType.BOUND_CAPSULE

    capsule_bpy = bound_to_bpy_object(
        capsule_xml, name=SOLLUMZ_UI_NAMES[sollum_type])
    capsule_bpy.sollum_type = sollum_type

    apply_bound_child_properties(capsule_bpy, capsule_xml)

    capsule_bpy.bound_length = capsule_xml.box_max.z - capsule_xml.box_min.z
    capsule_bpy.bound_radius = capsule_xml.sphere_radius

    return capsule_bpy


def bound_cylinder_to_bpy(cylinder_xml: BoundCylinder) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml BoundCylinder object."""
    sollum_type = SollumType.BOUND_CYLINDER

    cylinder_bpy = bound_to_bpy_object(
        cylinder_xml, name=SOLLUMZ_UI_NAMES[sollum_type])
    cylinder_bpy.sollum_type = sollum_type

    apply_bound_child_properties(cylinder_bpy, cylinder_xml)

    extents = cylinder_xml.box_max - cylinder_xml.box_min
    cylinder_bpy.bound_length = extents.y
    cylinder_bpy.bound_radius = extents.x * 0.5

    return cylinder_bpy


def bound_disc_to_bpy(disc_xml: BoundDisc) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml BoundDisc object."""
    sollum_type = SollumType.BOUND_DISC

    disc_bpy = bound_to_bpy_object(
        disc_xml, name=SOLLUMZ_UI_NAMES[sollum_type])
    disc_bpy.sollum_type = sollum_type

    apply_bound_child_properties(disc_bpy, disc_xml)

    disc_bpy.bound_radius = disc_xml.sphere_radius
    # TODO: Why is create_disc used?
    create_disc(disc_bpy.data, disc_xml.sphere_radius, disc_bpy.margin * 2)

    return disc_bpy


def material_xml_to_bpy_material(material_xml: MaterialItem) -> bpy.types.Material:
    """Create a bpy material from a cwxml material"""
    material_bpy = create_collision_material_from_index(material_xml.type)
    # Assign collision properties
    material_bpy.collision_properties.procedural_id = material_xml.procedural_id
    material_bpy.collision_properties.room_id = material_xml.room_id
    material_bpy.collision_properties.ped_density = material_xml.ped_density
    material_bpy.collision_properties.material_color_index = material_xml.material_color_index

    # Assign flags
    # TODO: Use same method as assign_composite_flags
    for flag_name in CollisionMatFlags.__annotations__.keys():
        if f"FLAG_{flag_name.upper()}" in material_xml.flags:
            setattr(material_bpy.collision_flags, flag_name, True)

    return material_bpy


def create_bound_polygon_mesh(vertices: list[Vector], triangles: list[Triangle], materials: list[bpy.types.Material], damaged: bool = False) -> bpy.types.Object:
    """Create a bpy mesh object that contains all triangle polygons from a cwxml BoundGeometry."""
    # TODO: Rename enum SollumType.BOUND_POLY_TRIANGLE2 to SollumType.BOUND_POLY_TRIANGLE_DAMAGED
    sollum_type = SollumType.BOUND_POLY_TRIANGLE if not damaged else SollumType.BOUND_POLY_TRIANGLE2

    name = SOLLUMZ_UI_NAMES[sollum_type]
    mesh = bpy.data.meshes.new(name)
    bpy_object = bpy.data.objects.new(name, mesh)
    bpy_object.sollum_type = sollum_type

    # Total vertices of mesh in order of faces
    mesh_vertices = []
    # Vertex indices of triangles [v1, v2, v3]
    faces = []
    # Index of material for each triangle
    material_indices = []

    for triangle in triangles:
        face_indices = []
        triangle_vertices = [vertices[triangle.v1],
                             vertices[triangle.v2], vertices[triangle.v3]]

        for vertex in triangle_vertices:
            if not vertex in mesh_vertices:
                mesh_vertices.append(vertex)
            face_indices.append(mesh_vertices.index(vertex))

        faces.append(face_indices)

        material_indices.append(triangle.material_index)

    try:
        mesh.from_pydata(vertices=mesh_vertices, edges=[], faces=faces)
    except:
        # TODO: Message dialog
        print(
            f"Error encountered while creating a {name} from XML vertex data:\n{format_exc()}")
        return bpy_object

    # TODO: Are extra materials added?
    for material in materials:
        mesh.materials.append(material)

    for index, triangle in mesh.polygons.items():
        if 0 <= index < len(material_indices):
            triangle.material_index = material_indices[index]

    bpy.context.collection.objects.link(bpy_object)

    return bpy_object


def polygon_bound_to_bpy(polygon_xml: Polygon, sollum_type: SollumType, materials: list[bpy.types.Material]):
    """Create a bpy mesh object from a generic cwxml bound polygon."""
    name = SOLLUMZ_UI_NAMES[sollum_type]
    mesh = bpy.data.meshes.new(name)
    if 0 <= polygon_xml.material_index < len(materials):
        mesh.materials.append(materials[polygon_xml.material_index])

    polygon_bpy = bpy.data.objects.new(name, mesh)
    polygon_bpy.sollum_type = sollum_type

    bpy.context.collection.objects.link(polygon_bpy)

    return polygon_bpy


def polygon_box_to_bpy(box_xml: Box, vertices: list[Vector], materials: list[bpy.types.Material]) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml bound polygon box."""
    box_bpy = polygon_bound_to_bpy(
        box_xml, SollumType.BOUND_POLY_BOX, materials)

    # Four opposing corners of the box
    v1 = vertices[box_xml.v1]
    v2 = vertices[box_xml.v2]
    v3 = vertices[box_xml.v3]
    v4 = vertices[box_xml.v4]
    center = (v1 + v2 + v3 + v4) * 0.25

    # Get edges from the 4 opposing corners of the box
    a1 = ((v3 + v4) - (v1 + v2)) * 0.5
    v2 = v1 + a1
    v3 = v3 - a1
    v4 = v4 - a1

    minedge = Vector((0.0001, 0.0001, 0.0001))
    edge1 = max(v2 - v1, minedge)
    edge2 = max(v3 - v1, minedge)
    edge3 = max((v4 - v1), minedge)

    # Order edges
    s1 = False
    s2 = False
    s3 = False
    if edge2.length > edge1.length:
        t1 = edge1
        edge1 = edge2
        edge2 = t1
        s1 = True
    if edge3.length > edge1.length:
        t1 = edge1
        edge1 = edge3
        edge3 = t1
        s2 = True
    if edge3.length > edge2.length:
        t1 = edge2
        edge2 = edge3
        edge3 = t1
        s3 = True

    # Ensure all edge vectors are perpendicular to each other
    b1 = edge1.normalized()
    b2 = edge2.normalized()
    b3 = b1.cross(b2).normalized()
    b2 = b1.cross(b3).normalized()
    edge2 = b2 * edge2.dot(b2)
    edge3 = b3 * edge3.dot(b3)

    # Unswap edges
    if s3 == True:
        t1 = edge2
        edge2 = edge3
        edge3 = t1
    if s2 == True:
        t1 = edge1
        edge1 = edge3
        edge3 = t1
    if s1 == True:
        t1 = edge1
        edge1 = edge2
        edge2 = t1

    # Construct transformation matrix using edge vectors
    mat = Matrix()
    mat[0] = edge1.x, edge2.x, edge3.x, center.x
    mat[1] = edge1.y, edge2.y, edge3.y, center.y
    mat[2] = edge1.z, edge2.z, edge3.z, center.z

    # Create unit cube (1x1x1)
    create_box(box_bpy.data, size=1)
    box_bpy.matrix_basis = mat

    return box_bpy


def polygon_sphere_to_bpy(sphere_xml: Sphere, vertices: list[Vector], materials: list[bpy.types.Material]) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml polygon sphere."""
    sphere_bpy = polygon_bound_to_bpy(
        sphere_xml, SollumType.BOUND_POLY_SPHERE, materials)

    sphere_bpy.bound_radius = sphere_xml.radius
    sphere_bpy.location = vertices[sphere_xml.v]

    return sphere_bpy


def polygon_capsule_to_bpy(capsule_xml: Capsule, vertices: list[Vector], materials: list[bpy.types.Material]) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml polygon capsule."""
    capsule_bpy = polygon_bound_to_bpy(
        capsule_xml, SollumType.BOUND_POLY_CAPSULE, materials)
    v1 = vertices[capsule_xml.v1]
    v2 = vertices[capsule_xml.v2]
    rot = get_direction_of_vectors(v1, v2)
    capsule_bpy.bound_radius = capsule_xml.radius * 2
    capsule_bpy.bound_length = (
        (v1 - v2).length + (capsule_xml.radius * 2)) / 2

    capsule_bpy.location = (v1 + v2) / 2
    capsule_bpy.rotation_euler = rot

    return capsule_bpy


def polygon_cylinder_to_bpy(cylinder_xml: Cylinder, vertices: list[Vector], materials: list[bpy.types.Material]) -> bpy.types.Object:
    """Create a bpy mesh object from a cwxml polygon cylinder."""
    cylinder_bpy = polygon_bound_to_bpy(
        cylinder_xml, SollumType.BOUND_POLY_CYLINDER, materials)
    v1 = vertices[cylinder_xml.v1]
    v2 = vertices[cylinder_xml.v2]
    rot = get_direction_of_vectors(v1, v2)
    cylinder_bpy.bound_radius = cylinder_xml.radius
    cylinder_bpy.bound_length = get_distance_of_vectors(v1, v2)
    cylinder_bpy.matrix_world = Matrix()

    cylinder_bpy.location = (v1 + v2) / 2
    cylinder_bpy.rotation_euler = rot

    return cylinder_bpy


# TODO: Needs testing after refactor
def bound_geometry_to_bpy(geometry_xml: Union[BoundGeometryBVH, BoundGeometry], bvh=True) -> bpy.types.Object:
    """Create a bpy empty containing the Bound Polygons as bpy mesh objects from a cwxml BoundGeometry/BoundGeometryBVH object."""
    geometry_bpy = bound_to_bpy_empty(
        geometry_xml, SOLLUMZ_UI_NAMES[SollumType.BOUND_GEOMETRYBVH if bvh else SollumType.BOUND_GEOMETRY])
    apply_bound_child_properties(geometry_bpy, geometry_xml)

    # Create materials
    materials = []
    for material_xml in geometry_xml.materials:
        materials.append(material_xml_to_bpy_material(material_xml))

    triangle_polygons = [
        polygon for polygon in geometry_xml.polygons if type(polygon) == Triangle]

    # Create a single mesh object from all the bound polygon triangles
    bound_polygon_mesh = create_bound_polygon_mesh(
        vertices=geometry_xml.vertices, triangles=triangle_polygons, materials=materials)
    bound_polygon_mesh.parent = geometry_bpy

    # BoundGeometry collisions can have a second vertex list that represents the vertices in a "damaged" state
    bound_polygon_mesh_damaged = None
    if not bvh:
        bound_polygon_mesh_damaged = create_bound_polygon_mesh(
            vertices=geometry_xml.vertices_2, triangles=triangle_polygons, materials=materials)
        bound_polygon_mesh_damaged.parent = geometry_bpy

    # Apply vertex colors to the mesh
    if len(geometry_xml.vertex_colors) > 0:
        create_vertexcolor_layer(mesh=bound_polygon_mesh.data, num=0,
                                 name="Color Layer", colors=geometry_xml.vertex_colors)

    for polygon in geometry_xml.polygons:
        polygon_bpy = None
        # TODO: Rename polygon classes and find way to avoid if statements
        if type(polygon) == Box:
            polygon_bpy = polygon_box_to_bpy(
                polygon, geometry_xml.vertices, bound_polygon_mesh.data.materials)
        elif type(polygon) == Sphere:
            polygon_bpy = polygon_sphere_to_bpy(
                polygon, geometry_xml.vertices, bound_polygon_mesh.data.materials)
        elif type(polygon) == Cylinder:
            polygon_bpy = polygon_cylinder_to_bpy(
                polygon, geometry_xml.vertices, bound_polygon_mesh.data.materials)
        elif type(polygon) == Capsule:
            polygon_bpy = polygon_capsule_to_bpy(
                polygon, geometry_xml.vertices, bound_polygon_mesh.data.materials)
        else:
            # TODO: Message dialog
            print(
                f"Failed to create objects for '{polygon_bpy}': {type(polygon).__name__} is an invalid bound type!")
            continue

        polygon_bpy.parent = geometry_bpy

    # TODO: Investigate
    # GeometryBVHs that have an UnkType of 2 have the triangle vertices offsetted by GeometryCenter.
    # Otherwise, the entire BVH is offsetted.
    if geometry_xml.unk_type == 2:
        bound_polygon_mesh.location = geometry_xml.geometry_center
        if bound_polygon_mesh_damaged is not None:
            bound_polygon_mesh_damaged.location += geometry_xml.geometry_center
    else:
        geometry_bpy.location += geometry_xml.geometry_center

    return geometry_bpy


def composite_to_bpy(composite_xml: BoundComposite, name: str) -> bpy.types.Object:
    """Create a bpy empty object from a cwxml BoundComposite object."""
    # Create composite bpy empty
    composite_bpy = bound_to_bpy_empty(composite_xml, name)

    # Create composite children
    for i, child in enumerate(composite_xml.children):
        child_bpy = None

        # TODO: Figure out how to do this without all the if statements
        if isinstance(child, BoundBox):
            child_bpy = bound_box_to_bpy(child)
        elif isinstance(child, BoundSphere):
            child_bpy = bound_sphere_to_bpy(child)
        elif isinstance(child, BoundCylinder):
            child_bpy = bound_cylinder_to_bpy(child)
        elif isinstance(child, BoundCapsule):
            child_bpy = bound_capsule_to_bpy(child)
        elif isinstance(child, BoundDisc):
            child_bpy = bound_disc_to_bpy(child)
        elif isinstance(child, BoundCloth):
            # NOTE: Not yet implemented
            child_bpy = bound_to_bpy_empty(child)
        elif isinstance(child, BoundGeometryBVH):
            child_bpy = bound_geometry_to_bpy(child)
        elif isinstance(child, BoundGeometry):
            child_bpy = bound_geometry_to_bpy(child)
        else:
            print(
                f"Failed to create objects for '{child}': {type(child).__name__} is an invalid bound type!")
            continue

        # TODO: Alternative method? If not, document this
        child_bpy.creation_index = i
        child_bpy.parent = composite_bpy

    return composite_bpy


def import_ybn(filepath: str) -> bpy.types.Object:
    """Import collisions into the scene from a .ybn.xml file."""
    ybn_xml = YBN.from_xml_file(filepath)

    if ybn_xml.composite is None:
        # TODO: Show message dialog
        print("No composite present in file")
        return

    return composite_to_bpy(composite_xml=ybn_xml.composite, name=os.path.basename(
        filepath.replace(YBN.file_extension, '')))