#include <assert.h>
#include <stdio.h>
#include <sys/stat.h>
#include <stdint.h>

#include <assimp/Importer.hpp>
#include <assimp/postprocess.h>
#include <assimp/scene.h>
#include <assimp/DefaultLogger.hpp>
#include <assimp/LogStream.hpp>

static void usage();
static void file_write_error_check(FILE *file, const void *data, size_t size);
static void file_write_u32_error_check(FILE *file, unsigned int value);
static void process_node(const aiScene *scene, const aiNode *node, FILE *output_file, const std::string &source_dir);

int main(int argc, char **argv) {
    if(argc != 3)
        usage();

    const char *source_filepath = argv[1];
    const char *destination_filepath = argv[2];

    Assimp::Logger::LogSeverity severity = Assimp::Logger::VERBOSE;
	Assimp::DefaultLogger::create("", severity, aiDefaultLogStream_STDERR);

    Assimp::Importer importer;
    const aiScene *scene = importer.ReadFile(source_filepath, aiProcessPreset_TargetRealtime_Quality | aiProcess_Triangulate | aiProcess_FlipUVs);
    if(!scene || scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode) {
        fprintf(stderr, "Failed to load file: %s, reason: %s\n", source_filepath, importer.GetErrorString());
        exit(2);
    }

    FILE *destination_file = fopen(destination_filepath, "wb");
    if(!destination_file) {
        perror(destination_filepath);
        exit(3);
    }
    file_write_u32_error_check(destination_file, 0x036144AF);  // magic number
    file_write_u32_error_check(destination_file, 0x00000001);  // version

    std::string source_dir = source_filepath;
    source_dir = source_dir.substr(0, source_dir.find_last_of("/"));

    process_node(scene, scene->mRootNode, destination_file, source_dir);
    fclose(destination_file);
    Assimp::DefaultLogger::kill();
    return 0;
}

void usage() {
    fprintf(stderr, "usage: amalgine-model-converter <source> <destination>\n");
    exit(1);
}

void file_write_error_check(FILE *file, const void *data, size_t size) {
    if(fwrite(data, 1, size, file) != size) {
        fprintf(stderr, "Error: failed to write %ld bytes to output file, error: %s\n", size, strerror(errno));
        exit(4);
    }
}

void file_write_u32_error_check(FILE *file, uint32_t value) {
    const size_t size = sizeof(value);
    if(fwrite(&value, 1, size, file) != size) {
        fprintf(stderr, "Error: failed to write %ld bytes to output file, error: %s\n", size, strerror(errno));
        exit(4);
    }
}

static bool is_file(const char *filename) {
    struct stat filestat;
    return stat(filename, &filestat) == 0 && S_ISREG(filestat.st_mode);
}

static void process_mesh(const aiScene *scene, const aiMesh *mesh, FILE *output_file, const std::string &source_dir) {
    // TODO: It seems like sometimes the vertice count doesn't match triangles.. excluding the last vertices in the list
    unsigned int vertice_overflow = mesh->mNumVertices % 3;
    unsigned int num_vertices = mesh->mNumVertices - vertice_overflow;
    assert(num_vertices % 3 == 0);
    file_write_u32_error_check(output_file, num_vertices);
    file_write_error_check(output_file, mesh->mVertices, num_vertices * sizeof(aiVector3D));
    file_write_error_check(output_file, mesh->mNormals, num_vertices * sizeof(aiVector3D));

    file_write_u32_error_check(output_file, mesh->mTextureCoords[0] ? 1 : 0);
    if(mesh->mTextureCoords[0]) {
        for(unsigned int i = 0; i < num_vertices; ++i) {
            file_write_error_check(output_file, &mesh->mTextureCoords[0][i].x, sizeof(float));
            file_write_error_check(output_file, &mesh->mTextureCoords[0][i].y, sizeof(float));
        }
    }

    file_write_u32_error_check(output_file, mesh->mNumFaces);
    for(unsigned int i = 0; i < mesh->mNumFaces; ++i) {
        // Should always be 3 since the file was loaded with aiProcess_Triangulate
        assert(mesh->mFaces[i].mNumIndices == 3);
        file_write_error_check(output_file, mesh->mFaces[i].mIndices, 3 * sizeof(unsigned int));
    }

    file_write_u32_error_check(output_file, mesh->mMaterialIndex < scene->mNumMaterials);
    if(mesh->mMaterialIndex < scene->mNumMaterials) {
        const unsigned int num_texture_types = 3;
        const aiTextureType texture_types[num_texture_types] = {
            aiTextureType_DIFFUSE,
            aiTextureType_SPECULAR,
            aiTextureType_AMBIENT
        };

        const aiMaterial *material = scene->mMaterials[mesh->mMaterialIndex];
        for(unsigned int i = 0; i < num_texture_types; ++i) {
            unsigned int num_textures_for_type = material->GetTextureCount(texture_types[i]);
            file_write_u32_error_check(output_file, num_textures_for_type);
            for(unsigned int j = 0; j < num_textures_for_type; ++j) {
                aiString texture_name;
                if(material->GetTexture(texture_types[i], j, &texture_name) == aiReturn_SUCCESS) {
                    std::string texture_path = source_dir + "/";
                    texture_path.append(texture_name.data, texture_name.length);
                    // This is needed because it's common for models to use a directory called textures or Textures
                    // and not specify it in the model file
                    if(!is_file(texture_path.c_str())) {
                        //fprintf(stderr, "Warning: texture not found: %s\n", texture_path.c_str());
                        texture_path = source_dir + "/Textures/";
                        texture_path.append(texture_name.data, texture_name.length);
                        if(!is_file(texture_path.c_str())) {
                            //fprintf(stderr, "Warning: texture not found: %s\n", texture_path.c_str());
                            texture_path = source_dir + "/textures/";
                            texture_path.append(texture_name.data, texture_name.length);
                            if(!is_file(texture_path.c_str())) {
                                fprintf(stderr, "Warning: texture not found: %s\n", texture_name.C_Str());
                                file_write_error_check(output_file, "\0", 1);
                                continue;
                            }
                        }
                    }

                    size_t dir_prefix_size = source_dir.size() + 1;
                    // +1 to include null-terminate character
                    file_write_error_check(output_file, texture_path.data() + dir_prefix_size, texture_path.size() + 1 - dir_prefix_size);
                } else {
                    file_write_error_check(output_file, "\0", 1);
                }
            }
        }
    }
}

void process_node(const aiScene *scene, const aiNode *node, FILE *output_file, const std::string &source_dir) {
    for(unsigned int i = 0; i < node->mNumMeshes; ++i) {
        const aiMesh *mesh = scene->mMeshes[node->mMeshes[i]];
        process_mesh(scene, mesh, output_file, source_dir);
    }

    for(unsigned int i = 0; i < node->mNumChildren; ++i) {
        process_node(scene, node->mChildren[i], output_file, source_dir);
    }
}