#include <destrum/Components/MeshRendererComponent.h>
#include <destrum/Assets/AssetReference.h>
#include <destrum/Graphics/GfxDevice.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 <stdexcept>

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

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

void MeshRendererComponent::Update(float) {
}

nlohmann::json MeshRendererComponent::Serialize() const {
    return {
        {"meshId", meshID},
        {"materialId", materialID},
        {"meshKey", meshKey},
        {"materialKey", materialKey}
    };
}

void MeshRendererComponent::Deserialize(const nlohmann::json& data) {
    if (data.contains("meshId")) {
        meshID = data.at("meshId").get<MeshID>();
    }
    if (data.contains("materialId")) {
        materialID = data.at("materialId").get<MaterialID>();
    }
    if (data.contains("meshKey")) {
        meshKey = data.at("meshKey").get<std::string>();
    }
    if (data.contains("materialKey")) {
        materialKey = data.at("materialKey").get<std::string>();
    }
}

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);
        if (!asset) {
            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) {
            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) {
            throw std::runtime_error("Material resource not found: " + materialKey);
        }
        materialID = *resolved;
        materialKey = resources->materials().getMaterialKey(*resolved);
    }

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