//
// Created by atdunbg on 2026/9/21.
//

#include <Engine/Renderer/Model.h>
#include <Engine/Renderer/Material.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Vulkan/MemoryAllocator.h>
#include <Engine/Vulkan/DescriptorManager.h>
#include <Engine/Vulkan/Texture.h>
#include <Engine/Core/Log.h>

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

#include <filesystem>
#include <algorithm>

namespace Engine {

Model::Model(Device& device, MemoryAllocator& allocator,
             DescriptorManager& descriptorManager,
             vk::DescriptorSetLayout textureLayout,
             const std::string& path)
    : m_Device(device)
    , m_Allocator(allocator)
    , m_DescriptorManager(descriptorManager)
    , m_TextureLayout(textureLayout)
{
    InitDefaultTextures();

    LoadModel(path);
}

Model::~Model() = default;

void Model::InitDefaultTextures()
{
    // 各材质槽位的兜底纹理（1x1，开销可忽略）：
    //   albedo 白 / MR 平参数（roughness 0.5, metallic 0）/ emissive 黑
    m_DefaultTextures[Material::Albedo] =
        Texture::CreateSolidColor(m_Device, m_Allocator, 255, 255, 255);
    m_DefaultTextures[Material::MetallicRoughness] =
        Texture::CreateSolidColor(m_Device, m_Allocator, 0, 128, 0, false);
    m_DefaultTextures[Material::Normal] =
        Texture::CreateSolidColor(m_Device, m_Allocator, 128, 128, 255, false);
    m_DefaultTextures[Material::Emissive] =
        Texture::CreateSolidColor(m_Device, m_Allocator, 0, 0, 0);
}

Model::Model(Device& device, MemoryAllocator& allocator,
             DescriptorManager& descriptorManager,
             vk::DescriptorSetLayout textureLayout,
             ProceduralTag, f32 halfSize, const std::array<u8, 3>& color)
    : m_Device(device)
    , m_Allocator(allocator)
    , m_DescriptorManager(descriptorManager)
    , m_TextureLayout(textureLayout)
{
    // 纯色 1x1 贴图（albedo 兜底）
    InitDefaultTextures();
    auto tex = Texture::CreateSolidColor(device, allocator, color[0], color[1], color[2]);

    // XZ 平面，法线 +Y，从上往下看 CCW 绕序
    const f32 h = halfSize;
    std::vector<Vertex> vertices = {
        {{-h, 0.0f, -h}, {0, 1, 0}, {0.0f, 0.0f}},
        {{-h, 0.0f,  h}, {0, 1, 0}, {0.0f, 0.0f}},
        {{ h, 0.0f,  h}, {0, 1, 0}, {0.0f, 0.0f}},
        {{ h, 0.0f, -h}, {0, 1, 0}, {0.0f, 0.0f}},
    };
    const std::vector<u32> indices = {0, 1, 2, 0, 2, 3};

    // 材质/纹理生命周期随模型（缓存持有）
    std::string paths[Material::kSlotCount] = {"", "", "", ""};
    Material* mat = GetOrCreateMaterial(paths, MaterialParams{});   // 全默认材质

    m_Meshes.push_back(CreateScope<Mesh>(
        m_Device, m_Allocator, vertices, indices, mat));
}

Scope<Model> Model::CreatePlane(Device& device, MemoryAllocator& allocator,
                                DescriptorManager& descriptorManager,
                                vk::DescriptorSetLayout textureLayout,
                                f32 halfSize, u8 r, u8 g, u8 b)
{
    return Scope<Model>(new Model(device, allocator, descriptorManager,
                                  textureLayout, ProceduralTag{},
                                  halfSize, std::array<u8, 3>{r, g, b}));
}

void Model::LoadModel(const std::string& path)
{
    m_ModelPath = path;

    // UV 翻转按格式区分：
    //   OBJ/FBX   —— v 原点在底部，需翻到左上惯例（配合"v=0=图像顶"采样）
    //   glTF/GLB  —— UV 本就是左上原点（v=0=图像顶），翻转会把所有
    //                贴图上下颠倒（法线/MR 全反 = 塑料感元凶）
    const std::string lower = [&] {
        std::string s = path;
        std::transform(s.begin(), s.end(), s.begin(),
                       [](unsigned char c) { return std::tolower(c); });
        return s;
    }();
    const bool isGltf = lower.size() >= 5 &&
        (lower.rfind(".gltf") == lower.size() - 5 ||
         lower.rfind(".glb") == lower.size() - 4);

    Assimp::Importer importer;
    const aiScene* scene = importer.ReadFile(path,
        aiProcess_Triangulate |
        (isGltf ? 0 : aiProcess_FlipUVs) |
        aiProcess_GenSmoothNormals |
        aiProcess_CalcTangentSpace |   // 法线贴图 TBN
        aiProcess_JoinIdenticalVertices);

    if (!scene || (scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE) || !scene->mRootNode) {
        ENGINE_ERROR("Assimp error: {}", importer.GetErrorString());
        return;
    }

    ProcessNode(scene->mRootNode, scene);
    ENGINE_INFO("Model loaded: {} ({} meshes, {} textures)",
                path, m_Meshes.size(), m_TextureCache.size());
}

void Model::ProcessNode(aiNode* node, const aiScene* scene)
{
    for (u32 i = 0; i < node->mNumMeshes; ++i) {
        aiMesh* mesh = scene->mMeshes[node->mMeshes[i]];
        m_Meshes.push_back(ProcessMesh(mesh, scene));
    }
    for (u32 i = 0; i < node->mNumChildren; ++i) {
        ProcessNode(node->mChildren[i], scene);
    }
}

Texture* Model::GetOrLoadTexture(const std::string& rawPath, bool srgb)
{
    // FBX/DCC 导出的路径常带反斜杠（Windows 语义），统一成正斜杠
    std::string path = rawPath;
    std::replace(path.begin(), path.end(), '\\', '/');

    // 缓存键带 sRGB 语义：同一路径可能同时被颜色槽与数据槽引用
    std::string key = (srgb ? "s:" : "l:") + path;
    auto it = m_TextureCache.find(key);
    if (it != m_TextureCache.end()) {
        return it->second.get();
    }

    // 文件不存在则用兜底纹理
    if (!std::filesystem::exists(path)) {
        ENGINE_WARN("Texture not found, using default: {}", path);
        return nullptr;
    }

    auto tex = CreateScope<Texture>(m_Device, m_Allocator, path, srgb);
    Texture* ptr = tex.get();
    m_TextureCache.emplace(key, std::move(tex));
    return ptr;
}

Material* Model::GetOrCreateMaterial(const std::string paths[Material::kSlotCount],
                                     const MaterialParams& params)
{
    const std::string key = Material::MakeCacheKey(paths, params);
    if (auto it = m_MaterialCache.find(key); it != m_MaterialCache.end()) {
        return it->second.get();
    }

    // 每槽位解析纹理：空/缺失折到该槽位默认纹理
    std::array<Texture*, Material::kSlotCount> textures{};
    // 槽位 sRGB 语义：颜色图 true，数据图 false
    const bool slotSrgb[Material::kSlotCount] = {true, false, false, true};

    for (u32 slot = 0; slot < Material::kSlotCount; ++slot) {
        Texture* tex = nullptr;
        if (!paths[slot].empty()) {
            tex = GetOrLoadTexture(paths[slot], slotSrgb[slot]);
        }
        textures[slot] = tex ? tex : m_DefaultTextures[slot].get();
    }

    auto mat = CreateScope<Material>(
        m_Device, m_Allocator, m_DescriptorManager, m_TextureLayout,
        textures, params);
    Material* ptr = mat.get();
    m_MaterialCache.emplace(key, std::move(mat));
    return ptr;
}

Scope<Mesh> Model::ProcessMesh(aiMesh* mesh, const aiScene* scene)
{
    std::vector<Vertex> vertices;
    std::vector<u32> indices;

    vertices.reserve(mesh->mNumVertices);

    for (u32 i = 0; i < mesh->mNumVertices; ++i) {
        Vertex vertex{};

        vertex.position = {
            mesh->mVertices[i].x,
            mesh->mVertices[i].y,
            mesh->mVertices[i].z
        };

        if (mesh->HasNormals()) {
            vertex.normal = {
                mesh->mNormals[i].x,
                mesh->mNormals[i].y,
                mesh->mNormals[i].z
            };
        }

        if (mesh->mTextureCoords[0]) {
            vertex.uv = {
                mesh->mTextureCoords[0][i].x,
                mesh->mTextureCoords[0][i].y
            };
        }

        // 切线（法线贴图 TBN；无切线数据置零，shader 侧退回几何法线）
        if (mesh->mTangents) {
            vertex.tangent = {
                mesh->mTangents[i].x,
                mesh->mTangents[i].y,
                mesh->mTangents[i].z
            };
        }

        vertices.push_back(vertex);
    }

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

    // 读取材质纹理路径（PBR metallic-roughness 槽位）+ 金属度/粗糙度因子
    std::string texPaths[Material::kSlotCount];
    MaterialParams params;

    if (mesh->mMaterialIndex >= 0) {
        aiMaterial* material = scene->mMaterials[mesh->mMaterialIndex];
        const std::filesystem::path modelDir =
            std::filesystem::path(m_ModelPath).parent_path();

        // glTF2 贴图在 assimp 5.x 的落点：MR 合图 → METALNESS(15)/
        // DIFFUSE_ROUGHNESS(16)（同图双槽）；baseColor → BASE_COLOR(12)。
        // 先到先得，数组顺序即优先级
        const std::pair<aiTextureType, u32> slots[] = {
            {aiTextureType_DIFFUSE,        Material::Albedo},
            {aiTextureType_BASE_COLOR,     Material::Albedo},
            {aiTextureType_NORMALS,        Material::Normal},
            {aiTextureType_METALNESS,      Material::MetallicRoughness},
            {aiTextureType_DIFFUSE_ROUGHNESS, Material::MetallicRoughness},
            {aiTextureType_UNKNOWN,        Material::MetallicRoughness},
            {aiTextureType_EMISSIVE,       Material::Emissive},
            {aiTextureType_EMISSION_COLOR, Material::Emissive},
        };

        for (const auto& [type, slot] : slots) {
            if (!texPaths[slot].empty()) continue;   // 高优先级已命中
            aiString path;
            if (material->GetTexture(type, 0, &path) == AI_SUCCESS &&
                path.length > 0) {
                texPaths[slot] = (modelDir / path.C_Str()).string();
            }
        }

        // 材质加载可见性：贴图槽缺失/路径解析错误会静默退化成默认
        // 纹理（金属全没 → 塑料感），这里留下诊断锚点
        ENGINE_INFO("Material '{}': albedo={} normal={} MR={} emissive={}",
                    material->GetName().C_Str(),
                    texPaths[Material::Albedo].empty() ? "<none>"
                        : texPaths[Material::Albedo],
                    texPaths[Material::Normal].empty() ? "<none>"
                        : texPaths[Material::Normal],
                    texPaths[Material::MetallicRoughness].empty() ? "<none>"
                        : texPaths[Material::MetallicRoughness],
                    texPaths[Material::Emissive].empty() ? "<none>"
                        : texPaths[Material::Emissive]);

        // PBR 因子（glTF2 材质属性；旧资产无此键 → 默认电介质 0/0.5）
        ai_real metallic = 0.0f, roughness = 0.5f;
        material->Get(AI_MATKEY_METALLIC_FACTOR, metallic);
        material->Get(AI_MATKEY_ROUGHNESS_FACTOR, roughness);
        params.factors = glm::vec4(static_cast<f32>(metallic),
                                   static_cast<f32>(roughness), 0.0f, 0.0f);
    }

    // 材质缓存：同贴图组合+因子 → 同一 Material（一个描述符集）
    Material* mat = GetOrCreateMaterial(texPaths, params);

    return CreateScope<Mesh>(
        m_Device, m_Allocator, vertices, indices, mat);
}

} // namespace Engine