//
// Created by lbmas on 4/29/2026.
//

#include <spdlog/spdlog.h>

#include "ModelManager.h"

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

#include <cstring>

#include "TextureManager.h"
#include "glm/ext/matrix_transform.hpp"

std::unordered_map<std::string, std::shared_ptr<Model>> ModelManager::models;

void process_ai_node(aiNode* node, const aiScene* scene, glm::mat4 transform, ModelNode& model);

void process_ai_mesh(const aiMesh* aiMesh, Mesh& mesh);

void process_ai_material(const aiMaterial* aiMat, Material& mat, const std::filesystem::path& modelDirectory);
void process_ai_material_diffuse(const aiMaterial *aiMat, Material& mat, const std::filesystem::path& modelDirectory);
void process_ai_material_specular(const aiMaterial *aiMat, Material& mat, const std::filesystem::path& modelDirectory);

auto zUpMatrix = glm::mat4(
    1.0f,  0.0f,  0.0f,  0.0f,  // Column 0
    0.0f,  0.0f,  1.0f,  0.0f,  // Column 1
    0.0f, -1.0f,  0.0f,  0.0f,  // Column 2
    0.0f,  0.0f,  0.0f,  1.0f   // Column 3
);

std::shared_ptr<Model> ModelManager::load_from_file(std::string_view _path, bool zUp)
{
    Assimp::Importer importer;

    const aiScene* scene = importer.ReadFile(_path.data(),
                                             aiProcess_Triangulate |
                                             aiProcess_GenSmoothNormals |
                                             aiProcess_RemoveRedundantMaterials);

    if (!scene || scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode)
    {
        spdlog::error("failed to load model {}: {}", _path, importer.GetErrorString());
        return nullptr;
    }

    auto model = std::make_shared<Model>();

    const std::filesystem::path modelPath = {_path};
    for (unsigned int i = 0; i < scene->mNumMaterials; i++) {
        aiMaterial* aiMat = scene->mMaterials[i];
        aiString matName;
        if (aiMat->Get(AI_MATKEY_NAME, matName) && strcmp(matName.data, "DefaultMaterial") == 0) {
            continue;
        }
        Material mat{};
        process_ai_material(aiMat, mat, modelPath.parent_path());
        model->materials.push_back(std::make_shared<Material>(mat));
    }

    for (unsigned int i = 0; i < scene->mNumMeshes; i++) {
        aiMesh* aiMesh = scene->mMeshes[i];
        Mesh mesh{};
        mesh.materialId = aiMesh->mMaterialIndex;
        process_ai_mesh(aiMesh, mesh);
        model->meshes.push_back(std::make_shared<Mesh>(mesh));
    }

    auto transform = glm::identity<glm::mat4>();
    if (zUp) {
        transform = zUpMatrix;
    }

    process_ai_node(scene->mRootNode, scene, transform, model->root);

    return model;
}

void process_ai_node(aiNode* aiNode, const aiScene* scene, glm::mat4 transform, ModelNode& node)
{
    const auto t = aiNode->mTransformation;
    transform = transform * glm::mat4(t.a1, t.a2, t.a3, t.a4, t.b1, t.b2, t.b3, t.b4, t.c1, t.c2, t.c3, t.c4, t.d1, t.d2, t.d3, t.d4);

    node.transform = transform;

    for (unsigned int i = 0; i < aiNode->mNumMeshes; i++) {
        node.meshIndices.push_back(aiNode->mMeshes[i]);
    }

    for (unsigned int i = 0; i < aiNode->mNumChildren; i++)
    {
        ModelNode newNode{};
        process_ai_node(aiNode->mChildren[i], scene, transform, newNode);
        node.children.push_back(newNode);
    }
}

void process_ai_mesh(const aiMesh* aiMesh, Mesh& mesh)
{
    auto positions = std::vector<float>(aiMesh->mNumVertices * 3);
    auto uvs = std::vector<float>(aiMesh->mNumVertices * 2);
    auto normals = std::vector<float>(aiMesh->mNumVertices * 3);
    auto indices = std::vector<unsigned int>();

    for (unsigned int i = 0; i < aiMesh->mNumVertices; i++)
    {
        auto position = glm::vec4(aiMesh->mVertices[i].x, aiMesh->mVertices[i].y, aiMesh->mVertices[i].z, 1.0f);
        positions[3 * i] = position.x;
        positions[3 * i + 1] = position.y;
        positions[3 * i + 2] = position.z;

        if (aiMesh->HasNormals())
        {
            aiVector3D norm = aiMesh->mNormals[i];
            normals[3 * i]     = norm.x;
            normals[3 * i + 1] = norm.y;
            normals[3 * i + 2] = norm.z;
        }
        else
        {
            normals[3 * i] = 0;
            normals[3 * i + 1] = 0;
            normals[3 * i + 2] = 0;
        }

        if (aiMesh->HasTextureCoords(0))
        {
            uvs[2 * i] = aiMesh->mTextureCoords[0][i].x;
            uvs[2 * i + 1] = aiMesh->mTextureCoords[0][i].y;
        }
        else
        {
            uvs[2 * i] = 0;
            uvs[2 * i + 1] = 0;
        }
    }

    for(unsigned int i = 0; i < aiMesh->mNumFaces; i++)
    {
        const aiFace face = aiMesh->mFaces[i];
        for(unsigned int j = 0; j < face.mNumIndices; j++)
            indices.push_back(face.mIndices[j]);
    }

    mesh = {positions, uvs, normals, indices};
}

void process_ai_material(const aiMaterial* aiMat, Material& mat, const std::filesystem::path& modelDirectory) {
    aiString matName;
    aiReturn ret = aiMat->Get(AI_MATKEY_NAME, matName);
    if (ret != AI_SUCCESS) {
        spdlog::error("Could not find material name");
    }

    // grab the phong stuff
    aiColor3D ambientColor {0.f, 0.f, 0.f};
    ret = aiMat->Get(AI_MATKEY_COLOR_DIFFUSE, ambientColor);
    if (ret != AI_SUCCESS) {
        spdlog::error("Could not find ambient color for material: {}", matName.C_Str());
        ambientColor = {.3f, .3f, .3f};
    }

    aiColor3D diffuseColor {1.f, 1.f, 1.f};
    ret = aiMat->Get(AI_MATKEY_COLOR_DIFFUSE, diffuseColor);
    if (ret != AI_SUCCESS) {
        spdlog::error("Could not find diffuse color for material: {}", matName.C_Str());
        diffuseColor = {1.f, 1.f, 1.f};
    }

    aiColor3D specularColor {1.f, 1.f, 1.f};
    ret = aiMat->Get(AI_MATKEY_COLOR_SPECULAR, specularColor);
    if (ret != AI_SUCCESS) {
        spdlog::error("Could not find specular color for material: {}", matName.C_Str());
        specularColor = {1.f, 1.f, 1.f};
    }

    float shininess = 32.f;
    ret = aiMat->Get(AI_MATKEY_SHININESS, shininess);
    if (ret != AI_SUCCESS) {
        spdlog::error("Could not find shininess for material: {}", matName.C_Str());
        shininess = 32.f;
    }

    // now grab the textures, first diffuse
    process_ai_material_diffuse(aiMat, mat, modelDirectory);
    process_ai_material_specular(aiMat, mat, modelDirectory);

    mat.name = matName.C_Str();
    mat.phong.ambient = glm::vec3{ambientColor.r, ambientColor.g, ambientColor.b};
    mat.phong.diffuse = glm::vec3{diffuseColor.r, diffuseColor.g, diffuseColor.b};
    mat.phong.specular = glm::vec3{specularColor.r, specularColor.g, specularColor.b};
    mat.phong.shininess = shininess;
}

void process_ai_material_diffuse(const aiMaterial *aiMat, Material& mat, const std::filesystem::path& modelDirectory) {
    aiString aiTextureName;
    aiMat->GetTexture(aiTextureType_DIFFUSE, 0, &aiTextureName);
    std::string textureName = aiTextureName.C_Str();

    if (textureName.empty()) {
        mat.diffuse = TextureManager::textures["default_diffuse"];
        return;
    }

    const auto textureFilename = std::filesystem::path(textureName).filename();
    const auto texturePath = modelDirectory / textureFilename;

    const std::string finalTextureName = textureFilename.string();
    if (TextureManager::textures.contains(finalTextureName)) {
        mat.diffuse = TextureManager::textures[finalTextureName];
        return;
    }

    if (std::filesystem::exists(texturePath)) {
        const auto texture = TextureManager::load_from_file(texturePath.string());
        if (!texture) {
            mat.diffuse = TextureManager::textures["default_diffuse"];
            return;
        }

        TextureManager::textures[finalTextureName] = texture;
        mat.diffuse = texture;
    } else {
        mat.diffuse = TextureManager::textures["default_diffuse"];
    }
}

void process_ai_material_specular(const aiMaterial *aiMat, Material& mat, const std::filesystem::path& modelDirectory) {
    aiString aiTextureName;
    aiMat->GetTexture(aiTextureType_SPECULAR, 0, &aiTextureName);
    std::string textureName = aiTextureName.C_Str();

    if (textureName.empty()) {
        mat.specular = TextureManager::textures["default_specular"];
        return;
    }

    const auto textureFilename = std::filesystem::path(textureName).filename();
    const auto texturePath = modelDirectory / textureFilename;

    const std::string finalTextureName = textureFilename.string();
    if (TextureManager::textures.contains(finalTextureName)) {
        mat.specular = TextureManager::textures[finalTextureName];
        return;
    }

    if (std::filesystem::exists(texturePath)) {
        const auto texture = TextureManager::load_from_file(texturePath.string());
        if (!texture) {
            mat.specular = TextureManager::textures["default_specular"];
            return;
        }

        TextureManager::textures[finalTextureName] = texture;
        mat.specular = texture;
    } else {
        mat.specular = TextureManager::textures["default_specular"];
    }
}