#include <destrum/Components/MeshRendererComponent.h>
#include <destrum/Assets/AssetReference.h>
#include <destrum/Graphics/GfxDevice.h>
#include <destrum/Graphics/Resources/PrimitiveMeshes.h>
#include <destrum/Graphics/Renderer.h>
#include <destrum/ObjectModel/Transform.h>

#include "destrum/Components/Animator.h"
#include "destrum/ObjectModel/GameObject.h"
#include "destrum/Util/GameState.h"

#include <imgui.h>

#include <filesystem>
#include <stdexcept>

#include <spdlog/spdlog.h>

MeshRendererComponent::MeshRendererComponent(GameObject& parent): Component(parent, "MeshRendererComponent") {

}

void MeshRendererComponent::Start() {
    Component::Start();
    ResolveReferences(ObjectMap{});
    if (GetGameObject()->GetComponent<Animator>() &&
        meshID != NULL_MESH_ID &&
        GameState::GetInstance().HasGfxDevice() &&
        GameState::GetInstance().HasRenderer()) {
        const auto& gfxDevice = GameState::GetInstance().Gfx();
        const auto& mesh = GameState::GetInstance().Renderer().getResources()->meshes().getMesh(meshID);

        m_skinnedMesh = std::make_unique<SkinnedMesh>();
        m_skinnedMesh->skinnedVertexBuffer = gfxDevice.createBuffer(
            mesh.numVertices * sizeof(CPUMesh::Vertex),
            VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
            VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT |
            VK_BUFFER_USAGE_VERTEX_BUFFER_BIT
        );

        const std::string bufferName = GetGameObject()->GetName() + " (skinned vertices)";
        vkutil::addDebugLabel(gfxDevice.getDevice(), m_skinnedMesh->skinnedVertexBuffer.buffer, bufferName.c_str());
    }
}

void MeshRendererComponent::Destroy()
{
    if (m_skinnedMesh && GameState::GetInstance().HasGfxDevice()) {
        auto& gfxDevice = GameState::GetInstance().Gfx();
        gfxDevice.destroyBuffer(m_skinnedMesh->skinnedVertexBuffer);
    }
    m_skinnedMesh.reset();
    Component::Destroy();
}

void MeshRendererComponent::Update(float) {
}

nlohmann::json MeshRendererComponent::Serialize() const {
    // Only keys are saved, mesh and material ids are cache slots that change between runs
    std::string savedMeshKey = meshKey;
    std::string savedMaterialKey = materialKey;
    if (GameState::GetInstance().HasRenderer()) {
        if (const auto* resources = GameState::GetInstance().Renderer().getResources()) {
            if (savedMeshKey.empty() && meshID != NULL_MESH_ID) {
                savedMeshKey = resources->meshes().getMeshKey(meshID);
            }
            if (savedMaterialKey.empty() && materialID != NULL_MATERIAL_ID) {
                savedMaterialKey = resources->materials().getMaterialKey(materialID);
            }
        }
    }
    if (savedMeshKey.empty() && meshID != NULL_MESH_ID) {
        throw std::runtime_error("MeshRendererComponent has a mesh without a key, it cannot be saved");
    }
    if (savedMaterialKey.empty() && materialID != NULL_MATERIAL_ID) {
        throw std::runtime_error("MeshRendererComponent has a material without a key, it cannot be saved");
    }

    return {
        {"meshKey", savedMeshKey},
        {"materialKey", savedMaterialKey}
    };
}

void MeshRendererComponent::Deserialize(const nlohmann::json& data) {
    meshKey = data.value("meshKey", std::string{});
    materialKey = data.value("materialKey", std::string{});

    // Older scenes also stored raw ids, only trust them when there is no key
    meshID = meshKey.empty() ? data.value("meshId", NULL_MESH_ID) : NULL_MESH_ID;
    materialID = materialKey.empty() ? data.value("materialId", NULL_MATERIAL_ID) : NULL_MATERIAL_ID;
}

void MeshRendererComponent::ResolveReferences(const ObjectMap&) {
    if (!GameState::GetInstance().HasRenderer()) {
        return;
    }

    RenderResources* resources = GameState::GetInstance().Renderer().getResources();
    if (resources == nullptr) {
        return;
    }

    const auto loadAssetForKey = [&](const std::string& key) {
        const auto asset = AssetReference::fromCacheKey(key);
        // Runtime-only keys such as "BulbMaterial#2" also contain a '#', they are not files
        if (!asset || (asset->path.find("://") == std::string::npos &&
                       !std::filesystem::exists(asset->path))) {
            return;
        }
        if (!GameState::GetInstance().HasGfxDevice()) {
            throw std::runtime_error(
                "Cannot load mesh asset without a graphics device: " + asset->path);
        }

        resources->loadModelAsset(GameState::GetInstance().Gfx(), asset->path);
    };

    if (!meshKey.empty()) {
        auto resolved = resources->meshes().findMeshByKey(meshKey);
        if (!resolved && PrimitiveMeshes::IsPrimitiveKey(meshKey)) {
            resolved = resources->meshes().findMeshByKey(PrimitiveMeshes::StripDuplicateSuffix(meshKey));
            if (!resolved) {
                auto primitive = PrimitiveMeshes::MakeFromKey(meshKey);
                if (!primitive) {
                    throw std::runtime_error("Invalid primitive mesh key: " + meshKey);
                }
                if (!GameState::GetInstance().HasGfxDevice()) {
                    throw std::runtime_error("Cannot build primitive mesh without a graphics device: " + meshKey);
                }
                resolved = resources->meshes().addMesh(GameState::GetInstance().Gfx(), *primitive);
            }
            meshKey = resources->meshes().getMeshKey(*resolved);
        }
        if (!resolved) {
            loadAssetForKey(meshKey);
            resolved = resources->meshes().findMeshByKey(meshKey);
        }
        if (!resolved) {
            throw std::runtime_error("Mesh resource not found: " + meshKey);
        }
        meshID = *resolved;
    }
    if (!materialKey.empty()) {
        auto resolved = resources->materials().findMaterialByKey(materialKey);
        if (!resolved) {
            loadAssetForKey(materialKey);
            resolved = resources->materials().findMaterialByKey(materialKey);
        }

        if (!resolved) {
            const auto meshAsset = AssetReference::fromCacheKey(meshKey);
            resolved = resources->materials().findMaterialByName(
                materialKey,
                meshAsset ? meshAsset->path : std::string_view{});
        }

        if (!resolved) {
            resolved = resources->materials().findMaterialByName(materialKey);
        }

        if (resolved) {
            materialID = *resolved;
            materialKey = resources->materials().getMaterialKey(*resolved);
        } else {
            // Keep the key so saving again does not lose which material was meant
            spdlog::warn("Material \"{}\" on \"{}\" not found, using the placeholder material",
                         materialKey, GetGameObject()->GetName());
            materialID = resources->materials().getPlaceholderMaterialId();
        }
    }

    if (meshID != NULL_MESH_ID && meshKey.empty()) {
        (void)resources->meshes().getMesh(meshID);
    }
    if (materialID != NULL_MATERIAL_ID && materialKey.empty()) {
        (void)resources->materials().getMaterial(materialID);
    }
}

void MeshRendererComponent::ImGuiInspector() {
    if (meshID == NULL_MESH_ID) {
        ImGui::TextDisabled("Mesh: <none>");
    } else {
        ImGui::Text("Mesh ID: %zu", meshID);
        ImGui::TextWrapped(
            "Mesh key: %s",
            meshKey.empty() ? "<not resolved>" : meshKey.c_str());
    }

    if (materialID == NULL_MATERIAL_ID) {
        ImGui::TextDisabled("Material: <none>");
    } else {
        ImGui::Text("Material ID: %u", materialID);
        ImGui::TextWrapped(
            "Material key: %s",
            materialKey.empty() ? "<not resolved>" : materialKey.c_str());
    }

    ImGui::Text(
        "Skinning: %s",
        m_skinnedMesh != nullptr ? "active" : "static");
}

void MeshRendererComponent::SetMeshID(MeshID id) {
    meshID = id;
    meshKey.clear();
    if (id != NULL_MESH_ID && GameState::GetInstance().HasRenderer()) {
        if (auto* resources = GameState::GetInstance().Renderer().getResources()) {
            meshKey = resources->meshes().getMeshKey(id);
        }
    }
}

void MeshRendererComponent::SetMaterialID(MaterialID id) {
    materialID = id;
    materialKey.clear();
    if (id != NULL_MATERIAL_ID && GameState::GetInstance().HasRenderer()) {
        if (auto* resources = GameState::GetInstance().Renderer().getResources()) {
            materialKey = resources->materials().getMaterialKey(id);
        }
    }
}

void MeshRendererComponent::Render(const RenderContext& ctx) {
    if (meshID == NULL_MESH_ID || materialID == NULL_MATERIAL_ID) return;

    if (auto* animator = GetGameObject()->GetComponent<Animator>(); animator && m_skinnedMesh) {
        const auto skeleton = GetGameObject()->GetComponent<Animator>()->getSkeleton();
        std::uint32_t frameIdx = GameState::GetInstance().Gfx().getCurrentFrameIndex();

        const std::size_t jointMatricesStartIndex = animator->uploadJointMatrices(
            ctx,
            *skeleton,
            frameIdx
        );

        ctx.renderer.drawSkinnedMesh(
            meshID,
            GetTransform().GetWorldMatrix(),
            materialID,
            m_skinnedMesh.get(),
            jointMatricesStartIndex
        );
    } else {
        ctx.renderer.drawMesh(meshID, GetTransform().GetWorldMatrix(), materialID);
    }
}