bl_info = {
    "name": "Game Exporter",
    "description": "Exports the game data.",
    "author": "Austin Morlan",
    "blender": (2, 93, 0),
    "location": "File > Export > Game (.bin)",
    "category": "Import-Export",
}

import bpy
import math
import re
import os
import numpy as np
import wave
from dataclasses import dataclass

from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty, EnumProperty
from bpy.types import Operator

@dataclass
class Vertex:
    position: np.array
    normal:   np.array
    texcoord: np.array

@dataclass
class Texture:
    data:     bytearray
    width:    int
    height:   int
    channels: int

@dataclass
class Material:
    name: str
    texture: Texture

@dataclass
class Mesh:
    name: str
    material_idx: int
    indices:  list
    vertices: list

@dataclass
class Collider:
    position: np.array
    extents: np.array

@dataclass
class Tile:
    x: int
    y: int

@dataclass
class Audio:
    name: str
    data: np.array
    length: int
    channel_count: int
    sample_rate: int
    volume: np.array
    is_looping: bool

@dataclass
class Transform:
    position: np.array
    rotation: np.array
    scale: np.array

@dataclass
class Object:
    transform: Transform

def load_mesh(materials, mesh):
    # Calculate triangles and split normals
    mesh.calc_loop_triangles()
    mesh.calc_normals_split()

    # Collect all of the indices so that we know how many unique vertices we have
    indices = []
    max_index = 0
    for tri in mesh.loop_triangles:
        loop_idx0 = tri.loops[0]
        loop_idx1 = tri.loops[1]
        loop_idx2 = tri.loops[2]

        indices.append(loop_idx0)
        indices.append(loop_idx1)
        indices.append(loop_idx2)

        max_index = max(max_index, max(loop_idx0, max(loop_idx1, loop_idx2)))

    indices = np.array(indices, dtype="uint16")

    vertices = [None] * (max_index+1)
    for tri in mesh.loop_triangles:
        vert_idx0 = tri.vertices[0]
        vert_idx1 = tri.vertices[1]
        vert_idx2 = tri.vertices[2]
        loop_idx0 = tri.loops[0]
        loop_idx1 = tri.loops[1]
        loop_idx2 = tri.loops[2]

        p0 = np.array([
            mesh.vertices[vert_idx0].co[0],
            mesh.vertices[vert_idx0].co[1],
            mesh.vertices[vert_idx0].co[2]], dtype="float32")

        p1 = np.array([
            mesh.vertices[vert_idx1].co[0],
            mesh.vertices[vert_idx1].co[1],
            mesh.vertices[vert_idx1].co[2]], dtype="float32")

        p2 = np.array([
            mesh.vertices[vert_idx2].co[0],
            mesh.vertices[vert_idx2].co[1],
            mesh.vertices[vert_idx2].co[2]], dtype="float32")

        n0 = np.array([
            tri.split_normals[0][0],
            tri.split_normals[0][1],
            tri.split_normals[0][2]], dtype="float32")

        n1 = np.array([
            tri.split_normals[1][0],
            tri.split_normals[1][1],
            tri.split_normals[1][2]], dtype="float32")

        n2 = np.array([
            tri.split_normals[2][0],
            tri.split_normals[2][1],
            tri.split_normals[2][2]], dtype="float32")

        t0 = np.array([
            mesh.uv_layers.active.data[loop_idx0].uv[0],
            mesh.uv_layers.active.data[loop_idx0].uv[1]], dtype="float32")

        t1 = np.array([
            mesh.uv_layers.active.data[loop_idx1].uv[0],
            mesh.uv_layers.active.data[loop_idx1].uv[1]], dtype="float32")

        t2 = np.array([
            mesh.uv_layers.active.data[loop_idx2].uv[0],
            mesh.uv_layers.active.data[loop_idx2].uv[1]], dtype="float32")

        v0 = Vertex(p0, n0, t0)
        v1 = Vertex(p1, n1, t1)
        v2 = Vertex(p2, n2, t2)

        vertices[loop_idx0] = v0
        vertices[loop_idx1] = v1
        vertices[loop_idx2] = v2

    material_idx = 65535
    if len(mesh.materials) > 0:
        material = mesh.materials[0]
        if material != None:
            material_name = mesh.materials[0].name

            for i, m in enumerate(materials):
                if m.name == material_name:
                    material_idx = i
                    break

    return Mesh(mesh.name, material_idx, indices, vertices)

def load_material(material):
    node = material.node_tree.nodes['Image Texture']

    data = bytearray([int(p * 255) for p in node.image.pixels])
    width = node.image.size[0]
    height = node.image.size[1]
    channels = node.image.channels

    texture = Texture(data, width, height, channels)

    return Material(material.name, texture)

def load_transform(item):
    position = np.array([round(item.location.x, 4), round(item.location.y, 4), round(item.location.z, 4)], dtype="float32")
    rotation = np.array([round(item.rotation_euler.x, 4), round(item.rotation_euler.y, 4), round(item.rotation_euler.z, 4)], dtype="float32")
    scale = np.array([round(item.scale.x, 4), round(item.scale.y, 4), round(item.scale.z, 4)], dtype="float32")

    return Transform(position, rotation, scale)

def load_collider(c):
    position = np.array([round(c.location.x, 4), round(c.location.y, 4), round(c.location.z, 4)], dtype="float32")
    extents = np.array([round(c.scale.x/2, 4), round(c.scale.y/2, 4), round(c.scale.z, 4)], dtype="float32")

    return Collider(position, extents)

def load_collision(collision):
    colliders = []

    for c in collision:
        colliders.append(load_collider(c))

    return colliders

def load_layout(layout):
    tiles = []

    for t in layout:
        # Offset origin (-1) and scale (/2)
        # Tiles are 2x2 in world but we want to view them as 1x1 for routing purposes
        x = (t.location.x - 1) / 2
        y = (t.location.y - 1) / 2

        tiles.append(Tile(x, y))

    return tiles

def load_audio(items):
    audio = []

    for i in items:
        path = bpy.path.abspath(i.fv_audio_source.file_path)
        wav = wave.open(path, 'rb')

        frame_count = wav.getnframes()
        channel_count = wav.getnchannels()
        sample_rate = wav.getframerate()
        sample_width = wav.getsampwidth()
        volume = np.array(i.fv_audio_source.volume, dtype="float32")
        is_looping = i.fv_audio_source.is_looping

        data_bytes = wav.readframes(frame_count)
        int_to_float = 1.0/32768.0

        data_int = np.frombuffer(data_bytes, dtype="int16")
        data_f32 = []

        for d in data_int:
            val = float(d*int_to_float)
            data_f32.append(val)

        data = np.array(data_f32, dtype="float32")

        audio.append(Audio(i.name, data, frame_count, channel_count, sample_rate, volume, is_looping))

    return audio

def traverse_children(obj):
    if len(obj.children) == 0:
        return []

    for child in obj.children:
        children = []
        children += traverse_children(child)

def load_node(node):
    children = traverse_children(node)

    for child in children:
        print(child)

def serialize_meshes(fp, meshes):
    mesh_count = len(meshes)
    fp.write(mesh_count.to_bytes(2, byteorder='little'))

    for m in meshes:
        name = m.name
        name_length = len(name)+1
        fp.write(name_length.to_bytes(2, byteorder='little'))
        fp.write(name.encode('UTF-8'))
        fp.write("\0".encode('UTF-8'))

        fp.write(m.material_idx.to_bytes(2, byteorder='little'))

        index_count = m.indices.size
        fp.write(index_count.to_bytes(2, byteorder='little'))
        fp.write(bytearray(m.indices))

        vertex_count = len(m.vertices)
        fp.write(vertex_count.to_bytes(2, byteorder='little'))
        for vertex in m.vertices:
            fp.write(bytearray(vertex.position))
            fp.write(bytearray(vertex.normal))
            fp.write(bytearray(vertex.texcoord))

def serialize_materials(fp, materials):
    material_count = len(materials)
    fp.write(material_count.to_bytes(2, byteorder='little'))

    for m in materials:
        name = m.name
        name_length = len(name)+1
        fp.write(name_length.to_bytes(2, byteorder='little'))
        fp.write(name.encode('UTF-8'))
        fp.write("\0".encode('UTF-8'))
        fp.write(m.texture.width.to_bytes(2, byteorder='little'))
        fp.write(m.texture.height.to_bytes(2, byteorder='little'))
        fp.write(m.texture.channels.to_bytes(1, byteorder='little'))
        fp.write(bytearray(m.texture.data))

def serialize_object(fp, obj):
    fp.write(bytearray(obj.transform.position))
    fp.write(bytearray(obj.transform.rotation))
    fp.write(bytearray(obj.transform.scale))

def serialize_objects(fp, objects):
    object_count = len(objects)
    fp.write(object_count.to_bytes(2, byteorder='little'))

    for obj in objects:
        serialize_object(fp, obj)

def serialize_collider(fp, collider):
    fp.write(bytearray(collider.position))
    fp.write(bytearray(collider.extents))

def serialize_colliders(fp, colliders):
    collider_count = len(colliders)
    fp.write(collider_count.to_bytes(2, byteorder='little'))

    for collider in colliders:
        serialize_collider(fp, collider)

def serialize_layout(fp, layout):
    tile_count = 32
    tiles = [0] * tile_count

    for tile in layout:
        x = int(tile.x)
        y = int(tile.y)

        value = tiles[y] | (1 << (31-x))
        tiles[y] = value

    fp.write(tile_count.to_bytes(1, byteorder='little'))
    for tile in tiles:
        fp.write(tile.to_bytes(4, byteorder='little'))

def serialize_audio(fp, audio):
    audio_count = len(audio)
    fp.write(audio_count.to_bytes(2, byteorder='little'))

    for a in audio:
        name = a.name
        name_length = len(name)+1
        fp.write(name_length.to_bytes(2, byteorder='little'))
        fp.write(name.encode('UTF-8'))
        fp.write("\0".encode('UTF-8'))
        fp.write(a.length.to_bytes(4, byteorder='little'))
        fp.write(a.sample_rate.to_bytes(4, byteorder='little'))
        fp.write(a.channel_count.to_bytes(1, byteorder='little'))
        fp.write(bytearray(a.volume))
        fp.write(a.is_looping.to_bytes(1, byteorder='little'))
        fp.write(bytearray(a.data))

def serialize_world(path):
    print(f'Serializing world')

    materials = []
    for m in bpy.data.materials:
        if m.name == "Dots Stroke":
            continue

        material = load_material(m)
        materials.append(material)

    meshes = []
    for m in bpy.data.meshes:
        mesh = load_mesh(materials, m)
        meshes.append(mesh)

    coins = []
    for o in bpy.data.collections['Coins'].objects:
        transform = load_transform(o)
        obj = Object(transform)
        coins.append(obj)

    eyeballs = []
    for o in bpy.data.collections['Eyeballs'].objects:
        transform = load_transform(o)
        obj = Object(transform)
        eyeballs.append(obj)

    pedestals = []
    for o in bpy.data.collections['Pedestals'].objects:
        transform = load_transform(o)
        obj = Object(transform)
        pedestals.append(obj)

    score_hundreds = Object(load_transform(bpy.data.objects["ScoreHundreds"]))
    score_tens = Object(load_transform(bpy.data.objects["ScoreTens"]))
    score_ones = Object(load_transform(bpy.data.objects["ScoreOnes"]))
    score_coin = Object(load_transform(bpy.data.objects["ScoreCoin"]))
    score_collider = load_collider(bpy.data.objects["ScoreCollider"])

    # Collision
    colliders = load_collision(bpy.data.collections['Walls'].objects)

    # Level layout
    layout = load_layout(bpy.data.collections['Layout'].objects)

    # Load audio
    music = load_audio(bpy.data.collections['Music'].objects)
    sfx = load_audio(bpy.data.collections['SFX'].objects)

    fp = open(path, "wb")
    fp.write("AJM".encode('UTF-8'))
    fp.write("\0".encode('UTF-8'))

    serialize_meshes(fp, meshes)
    serialize_materials(fp, materials)
    serialize_objects(fp, coins)
    serialize_objects(fp, eyeballs)
    serialize_objects(fp, pedestals)
    serialize_object(fp, score_hundreds)
    serialize_object(fp, score_tens)
    serialize_object(fp, score_ones)
    serialize_object(fp, score_coin)
    serialize_collider(fp, score_collider)
    serialize_colliders(fp, colliders)
    serialize_layout(fp, layout)
    serialize_audio(fp, music)
    serialize_audio(fp, sfx)

    fp.close()

def export(path):
    serialize_world(path)
    return {'FINISHED'}

class ExportGame(Operator, ExportHelper):
    bl_idname = "game_export.scene_text"
    bl_label = "Export Game (bin) File"
    filename_ext = ".bin"
    filter_glob: StringProperty(
        default="*.bin",
        options={'HIDDEN'},
        maxlen=255)

    def execute(self, context):
        return export(self.filepath)

def menu_func_export(self, context):
    self.layout.operator(ExportGame.bl_idname, text="Game (.bin) Custom")

def register():
    bpy.utils.register_class(ExportGame)
    bpy.types.TOPBAR_MT_file_export.append(menu_func_export)

def unregister():
    bpy.utils.unregister_class(ExportGame)
    bpy.types.TOPBAR_MT_file_export.remove(menu_func_export)

if __name__ == "__main__":
    register()