//
// 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 <meshoptimizer.h>

#include <filesystem>
#include <algorithm>

namespace Engine {

Model::Model(Device& device, MemoryAllocator& allocator,
             DescriptorManager& descriptorManager,
             vk::DescriptorSetLayout textureLayout,
             const std::string& path)
    : Model(device, allocator, descriptorManager, textureLayout)
{
    // 同步便捷路径 = 两段式串行（等价异步收割后的主线程消费）
    LoadModel(path);
}

// 空模型骨架：只建默认纹理；mesh/材质由 FromParsed 填充
Model::Model(Device& device, MemoryAllocator& allocator,
             DescriptorManager& descriptorManager,
             vk::DescriptorSetLayout textureLayout)
    : m_Device(device)
    , m_Allocator(allocator)
    , m_DescriptorManager(descriptorManager)
    , m_TextureLayout(textureLayout)
{
    InitDefaultTextures();
}

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);
    m_DefaultTextures[Material::Ao] =
        Texture::CreateSolidColor(m_Device, m_Allocator, 255, 255, 255, false);
}

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};
    const std::vector<std::vector<u32>> lodIndices = {indices};

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

    m_Meshes.push_back(CreateScope<Mesh>(
        m_Device, m_Allocator, vertices, lodIndices, 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}));
}

ParsedModel Model::ParseFromDisk(const std::string& path)
{
    ParsedModel out;
    out.path = 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;
    }();
    out.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 |
        (out.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 out;   // meshes 空 = 无效
    }

    const std::filesystem::path modelDir =
        std::filesystem::path(path).parent_path();
    ProcessNode(scene->mRootNode, scene, out.meshes, modelDir, out.isGltf);
    ENGINE_INFO("Model parsed: {} ({} meshes)", path, out.meshes.size());
    return out;
}

Scope<Model> Model::FromParsed(Device& device, MemoryAllocator& allocator,
                               DescriptorManager& descriptorManager,
                               vk::DescriptorSetLayout textureLayout,
                               ParsedModel&& parsed)
{
    if (!parsed.Valid()) {
        ENGINE_ERROR("Model::FromParsed: invalid parsed model ({})", parsed.path);
        return nullptr;
    }

    auto model = Scope<Model>(new Model(device, allocator,
                                        descriptorManager, textureLayout));
    model->m_ModelPath = parsed.path;
    model->m_IsGltf = parsed.isGltf;

    // 材质缓存去重 + 纹理上传 + mesh 缓冲（全部 GPU，主线程）
    for (auto& pm : parsed.meshes) {
        Material* mat = model->GetOrCreateMaterial(pm.texPaths.data(), pm.params);
        model->m_Meshes.push_back(CreateScope<Mesh>(
            device, allocator, pm.vertices, pm.lods, mat));
    }
    ENGINE_INFO("Model loaded: {} ({} meshes, {} textures)",
                parsed.path, model->m_Meshes.size(), model->m_TextureCache.size());
    return model;
}

void Model::LoadModel(const std::string& path)
{
    // 同步路径：解析 + GPU 构建串行（与 FromParsed 相同的消费逻辑）
    ParsedModel parsed = ParseFromDisk(path);
    if (!parsed.Valid()) {
        return;
    }
    m_ModelPath = parsed.path;
    m_IsGltf = parsed.isGltf;
    for (auto& pm : parsed.meshes) {
        Material* mat = GetOrCreateMaterial(pm.texPaths.data(), pm.params);
        m_Meshes.push_back(CreateScope<Mesh>(
            m_Device, m_Allocator, pm.vertices, pm.lods, mat));
    }
    ENGINE_INFO("Model loaded: {} ({} meshes, {} textures)",
                path, m_Meshes.size(), m_TextureCache.size());
}

void Model::ProcessNode(const aiNode* node, const aiScene* scene,
                        std::vector<ParsedMesh>& out,
                        const std::filesystem::path& modelDir, bool isGltf)
{
    for (u32 i = 0; i < node->mNumMeshes; ++i) {
        aiMesh* mesh = scene->mMeshes[node->mMeshes[i]];
        out.push_back(ProcessMesh(mesh, scene, modelDir, isGltf));
    }
    for (u32 i = 0; i < node->mNumChildren; ++i) {
        ProcessNode(node->mChildren[i], scene, out, modelDir, isGltf);
    }
}

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;
}

ParsedMesh Model::ProcessMesh(const aiMesh* mesh, const aiScene* scene,
                              const std::filesystem::path& modelDir,
                              bool isGltf)
{
    ParsedMesh out;
    std::vector<Vertex>& vertices = out.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) {
        const aiMaterial* material = scene->mMaterials[mesh->mMaterialIndex];

        // glTF2 贴图在 assimp 5.x 的落点：MR 合图 → METALNESS(15)/
        // DIFFUSE_ROUGHNESS(16)（同图双槽）；baseColor → BASE_COLOR(12)。
        // OBJ：map_Ns → SHININESS（独立 roughness 图）、map_Ka 常绑 AO。
        // 先到先得，数组顺序即优先级
        const std::pair<aiTextureType, u32> slots[] = {
            {aiTextureType_DIFFUSE,           Material::Albedo},
            {aiTextureType_BASE_COLOR,        Material::Albedo},
            {aiTextureType_NORMALS,           Material::Normal},
            {aiTextureType_HEIGHT,            Material::Normal},   // OBJ map_Bump 落点
            {aiTextureType_METALNESS,         Material::MetallicRoughness},
            {aiTextureType_DIFFUSE_ROUGHNESS, Material::MetallicRoughness},
            {aiTextureType_UNKNOWN,           Material::MetallicRoughness},
            {aiTextureType_SHININESS,         Material::MetallicRoughness},
            {aiTextureType_AMBIENT_OCCLUSION, Material::Ao},
            {aiTextureType_LIGHTMAP,          Material::Ao},
            {aiTextureType_AMBIENT,           Material::Ao},
            {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={} ao={}",
                    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],
                    texPaths[Material::Ao].empty() ? "<none>"
                        : texPaths[Material::Ao]);

        // 因子按工作流分支：
        //   glTF2：MR 合图（B=金属度语义），因子按声明（缺省 1/1）
        //   OBJ：specular 工作流无金属概念 → metallic=0；
        //     roughness 有独立贴图 → 因子 1（G 通道直达）；
        //     无贴图 → Blinn-Phong Ns 换算 R=sqrt(2/(Ns+2))，
        //     补偿默认 MR 贴图 G=0.5 → 因子 = 2R
        if (isGltf) {
            ai_real metallic = 1.0f, roughness = 1.0f;
            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);
        } else {
            f32 roughnessFactor = 1.0f;
            if (texPaths[Material::MetallicRoughness].empty()) {
                ai_real shininess = 10.0f;
                material->Get(AI_MATKEY_SHININESS, shininess);
                const f32 rough =
                    std::sqrt(2.0f / (static_cast<f32>(shininess) + 2.0f));
                roughnessFactor = 2.0f * glm::clamp(rough, 0.05f, 1.0f);
            }
            params.factors = glm::vec4(0.0f, roughnessFactor, 0.0f, 0.0f);
        }
    }

    // 材质创建延迟到 GPU 构建段（FromParsed）：解析期只留引用
    for (u32 slot = 0; slot < Material::kSlotCount; ++slot) {
        out.texPaths[slot] = std::move(texPaths[slot]);
    }
    out.params = params;

    // ---- LOD 链（解析期生成，worker 线程免费时间）：----
    // LOD0 = 原始；每级目标索引数减半，相对误差 1e-2。顶点共享，
    // 简化只重排索引（meshopt_simplify 从 LOD0 单次起点逐级向下）
    out.lods.push_back(std::move(indices));
    for (u32 lod = 1; lod < kMeshLodCount; ++lod) {
        const size_t target = out.lods[0].size() >> lod;
        // 太小的网格不简化：简化收益低于误差风险，直接复用上一级
        if (target < 96 || out.lods[0].size() < 192) {
            out.lods.push_back(out.lods.back());
            continue;
        }
        std::vector<u32> dest(out.lods[0].size());
        float error = 0.0f;
        const size_t count = meshopt_simplify(
            dest.data(),
            out.lods[0].data(), out.lods[0].size(),
            &out.vertices[0].position.x, out.vertices.size(), sizeof(Vertex),
            target, 1e-2f, 0, &error);
        if (count == 0) {
            out.lods.push_back(out.lods.back());   // 简化失败兜底
            continue;
        }
        dest.resize(count);
        out.lods.push_back(std::move(dest));
    }
    return out;
}

} // namespace Engine