import bpy

from .util import escape_name, mat2str, needs_def_mb, needs_xform_mb, ExportCancelled
from mathutils import Vector, Matrix

def make_object_data_cache(render_engine, depsgraph, ob, name):
    if ob.type == 'MESH':
        return Mesh(render_engine, depsgraph, ob, name)
    elif ob.type == 'LIGHT':
        if ob.data.type == 'POINT':
            return SphereLamp(render_engine, depsgraph, ob, name)
        elif ob.data.type == 'AREA':
            return RectLamp(render_engine, depsgraph, ob, name)
        elif ob.data.type == 'AREA':
            return RectLamp(render_engine, depsgraph, ob, name)

class Mesh:
    """ Holds data for a mesh to be exported.
    """
    def __init__(self, render_engine, depsgraph, ob, name):
        self.name = name
        self.material_name = None
        if len(ob.material_slots) >= 1 and ob.material_slots[0].material != None:
            self.material_name = ob.material_slots[0].material.name
        self.is_subdiv = ob.data.psychopath.is_subdivision_surface
        self.needs_mb = needs_def_mb(ob)
        self.time_meshes = []

    def take_sample(self, render_engine, depsgraph, ob, time):
        if len(self.time_meshes) == 0 or self.needs_mb:
            render_engine.update_stats("", "Psychopath: Collecting '{}' at time {}".format(self.name, time))
            self.time_meshes += [ob.to_mesh(depsgraph=depsgraph).copy()]
    
    def cleanup(self):
        for mesh in self.time_meshes:
            bpy.data.meshes.remove(mesh)

    def export(self, render_engine, w):
        render_engine.update_stats("", "Psychopath: Exporting %s" % self.name)

        if self.is_subdiv == False:
            # Exporting normal mesh
            w.write("MeshSurface $%s {\n" % self.name)
            w.indent()
        else:
            # Exporting subdivision surface cage
            w.write("SubdivisionSurface $%s {\n" % self.name)
            w.indent()

        # Material bindings.
        if self.material_name != None:
            w.write("SurfaceShaderBind [${}]\n".format(self.material_name))

        # Write vertices and (if it's smooth shaded) normals
        for ti in range(len(self.time_meshes)):
            w.write("Vertices [")
            w.write(" ".join([("%f" % i) for vert in self.time_meshes[ti].vertices for i in vert.co]), False)
            w.write("]\n", False)
            if self.time_meshes[0].polygons[0].use_smooth and self.is_subdiv == False:
                w.write("Normals [")
                w.write(" ".join([("%f" % i) for vert in self.time_meshes[ti].vertices for i in vert.normal]), False)
                w.write("]\n", False)

        # Write face vertex counts
        w.write("FaceVertCounts [")
        w.write(" ".join([("%d" % len(p.vertices)) for p in self.time_meshes[0].polygons]), False)
        w.write("]\n", False)

        # Write face vertex indices
        w.write("FaceVertIndices [")
        w.write(" ".join([("%d"%v) for p in self.time_meshes[0].polygons for v in p.vertices]), False)
        w.write("]\n", False)

        # MeshSurface/SubdivisionSurface section end
        w.unindent()
        w.write("}\n")


class SphereLamp:
    """ Holds data for a sphere light to be exported.
    """
    def __init__(self, render_engine, depsgraph, ob, name):
        self.name = name
        self.time_col = []
        self.time_rad = []

    def take_sample(self, render_engine, depsgraph, ob, time):
        render_engine.update_stats("", "Psychopath: Collecting '{}' at time {}".format(ob.name, time))

        if ob.data.psychopath.color_type == 'Rec709':
            self.time_col += [('Rec709', ob.data.color * ob.data.energy)]
        elif ob.data.psychopath.color_type == 'Blackbody':
            self.time_col += [('Blackbody', ob.data.psychopath.color_blackbody_temp, ob.data.energy)]
        elif ob.data.psychopath.color_type == 'ColorTemperature':
            self.time_col += [('ColorTemperature', ob.data.psychopath.color_blackbody_temp, ob.data.energy)]

        self.time_rad += [ob.data.shadow_soft_size]

    def cleanup(self):
        pass
    
    def export(self, render_engine, w):
        render_engine.update_stats("", "Psychopath: Exporting %s" % self.name)

        w.write("SphereLight $%s {\n" % self.name)
        w.indent()
        for col in self.time_col:
            if col[0] == 'Rec709':
                w.write("Color [rec709, %f %f %f]\n" % (col[1][0], col[1][1], col[1][2]))
            elif col[0] == 'Blackbody':
                w.write("Color [blackbody, %f %f]\n" % (col[1], col[2]))
            elif col[0] == 'ColorTemperature':
                w.write("Color [color_temperature, %f %f]\n" % (col[1], col[2]))
        for rad in self.time_rad:
            w.write("Radius [%f]\n" % rad)

        w.unindent()
        w.write("}\n")


class RectLamp:
    """ Holds data for a rectangular light to be exported.
    """
    def __init__(self, render_engine, depsgraph, ob, name):
        self.name = name
        self.time_col = []
        self.time_dim = []

    def take_sample(self, render_engine, depsgraph, ob, time):
        render_engine.update_stats("", "Psychopath: Collecting '{}' at time {}".format(self.name, time))

        if ob.data.psychopath.color_type == 'Rec709':
            self.time_col += [('Rec709', ob.data.color * ob.data.energy)]
        elif ob.data.psychopath.color_type == 'Blackbody':
            self.time_col += [('Blackbody', ob.data.psychopath.color_blackbody_temp, ob.data.energy)]
        elif ob.data.psychopath.color_type == 'ColorTemperature':
            self.time_col += [('ColorTemperature', ob.data.psychopath.color_blackbody_temp, ob.data.energy)]

        if ob.data.shape == 'RECTANGLE':
            self.time_dim += [(ob.data.size, ob.data.size_y)]
        else:
            self.time_dim += [(ob.data.size, ob.data.size)]

    def cleanup(self):
        pass
    
    def export(self, render_engine, w):
        render_engine.update_stats("", "Psychopath: Exporting %s" % self.name)

        w.write("RectangleLight $%s {\n" % self.name)
        w.indent()
        for col in self.time_col:
            if col[0] == 'Rec709':
                w.write("Color [rec709, %f %f %f]\n" % (col[1][0], col[1][1], col[1][2]))
            elif col[0] == 'Blackbody':
                w.write("Color [blackbody, %f %f]\n" % (col[1], col[2]))
            elif col[0] == 'ColorTemperature':
                w.write("Color [color_temperature, %f %f]\n" % (col[1], col[2]))
        for dim in self.time_dim:
            w.write("Dimensions [%f %f]\n" % dim)

        w.unindent()
        w.write("}\n")


class DistantDiskLamp:
    def __init__(self, render_engine, depsgraph, ob, name):
        self.name = name
        self.time_col = []
        self.time_dir = []
        self.time_rad = []

    def take_sample(self, render_engine, depsgraph, ob, time):
        render_engine.update_stats("", "Psychopath: Collecting '{}' at time {}".format(self.name, time))
        self.time_dir += [tuple(ob.matrix_world.to_3x3() @ Vector((0, 0, -1)))]

        if ob.data.psychopath.color_type == 'Rec709':
            self.time_col += [('Rec709', ob.data.color * ob.data.energy)]
        elif ob.data.psychopath.color_type == 'Blackbody':
            self.time_col += [('Blackbody', ob.data.psychopath.color_blackbody_temp, ob.data.energy)]
        elif ob.data.psychopath.color_type == 'ColorTemperature':
            self.time_col += [('ColorTemperature', ob.data.psychopath.color_blackbody_temp, ob.data.energy)]

        self.time_rad += [ob.data.shadow_soft_size]

    def cleanup(self):
        pass

    def export(self, render_engine, w):
        render_engine.update_stats("", "Psychopath: Exporting %s" % self.name)
        w.write("DistantDiskLight $%s {\n" % self.name)
        w.indent()
        for direc in self.time_dir:
            w.write("Direction [%f %f %f]\n" % (direc[0], direc[1], direc[2]))
        for col in self.time_col:
            if col[0] == 'Rec709':
                w.write("Color [rec709, %f %f %f]\n" % (col[1][0], col[1][1], col[1][2]))
            elif col[0] == 'Blackbody':
                w.write("Color [blackbody, %f %f]\n" % (col[1], col[2]))
            elif col[0] == 'ColorTemperature':
                w.write("Color [color_temperature, %f %f]\n" % (col[1], col[2]))
        for rad in self.time_rad:
            w.write("Radius [%f]\n" % rad)

        w.unindent()
        w.write("}\n")


# class Instance:
#     def __init__(self, render_engine, depsgraph, ob, data_name):
#         self.data_name = data_name
#         self.needs_mb = needs_xform_mb(self.ob)
#         self.time_xforms = []

#     def take_sample(self, render_engine, time, translation_offset):
#         if len(self.time_xforms) == 0 or self.needs_mb:
#             render_engine.update_stats("", "Psychopath: Collecting '{}' xforms at time {}".format(self.ob.name, time))
#             mat = self.ob.matrix_world.copy()
#             mat[0][3] += translation_offset[0]
#             mat[1][3] += translation_offset[1]
#             mat[2][3] += translation_offset[2]
#             self.time_xforms += [mat]

#     def export(self, render_engine, w):
#         render_engine.update_stats("", "Psychopath: Exporting %s" % self.ob.name)

#         w.write("Instance {\n")
#         w.indent()
#         w.write("Data [$%s]\n" % self.data_name)
#         for mat in self.time_xforms:
#             w.write("Transform [%s]\n" % mat2str(mat))
#         for ms in self.ob.material_slots:
#             if ms != None:
#                 w.write("SurfaceShaderBind [$%s]\n" % escape_name(ms.material.name))
#                 break
#         w.unindent()
#         w.write("}\n")