#include "Rendering/RenderResourceCache.h"

#include <XCEngine/RHI/RHIEnums.h>

namespace XCEngine {
namespace Rendering {

namespace {

RHI::Format ToRHITextureFormat(Resources::TextureFormat format) {
    switch (format) {
        case Resources::TextureFormat::RGBA8_UNORM:
        case Resources::TextureFormat::RGBA8_SRGB:
            return RHI::Format::R8G8B8A8_UNorm;
        default:
            return RHI::Format::Unknown;
    }
}

RHI::TextureType ToRHITextureType(Resources::TextureType textureType) {
    switch (textureType) {
        case Resources::TextureType::Texture2DArray:
            return RHI::TextureType::Texture2DArray;
        case Resources::TextureType::Texture3D:
            return RHI::TextureType::Texture3D;
        case Resources::TextureType::TextureCube:
            return RHI::TextureType::TextureCube;
        case Resources::TextureType::TextureCubeArray:
            return RHI::TextureType::TextureCubeArray;
        case Resources::TextureType::Texture2D:
        default:
            return RHI::TextureType::Texture2D;
    }
}

void ShutdownMesh(RenderResourceCache::CachedMesh& cachedMesh) {
    if (cachedMesh.vertexBufferView != nullptr) {
        cachedMesh.vertexBufferView->Shutdown();
        delete cachedMesh.vertexBufferView;
        cachedMesh.vertexBufferView = nullptr;
    }

    if (cachedMesh.indexBufferView != nullptr) {
        cachedMesh.indexBufferView->Shutdown();
        delete cachedMesh.indexBufferView;
        cachedMesh.indexBufferView = nullptr;
    }

    if (cachedMesh.vertexBuffer != nullptr) {
        cachedMesh.vertexBuffer->Shutdown();
        delete cachedMesh.vertexBuffer;
        cachedMesh.vertexBuffer = nullptr;
    }

    if (cachedMesh.indexBuffer != nullptr) {
        cachedMesh.indexBuffer->Shutdown();
        delete cachedMesh.indexBuffer;
        cachedMesh.indexBuffer = nullptr;
    }
}

void ShutdownTexture(RenderResourceCache::CachedTexture& cachedTexture) {
    if (cachedTexture.shaderResourceView != nullptr) {
        cachedTexture.shaderResourceView->Shutdown();
        delete cachedTexture.shaderResourceView;
        cachedTexture.shaderResourceView = nullptr;
    }

    if (cachedTexture.texture != nullptr) {
        cachedTexture.texture->Shutdown();
        delete cachedTexture.texture;
        cachedTexture.texture = nullptr;
    }
}

} // namespace

RenderResourceCache::~RenderResourceCache() {
    Shutdown();
}

void RenderResourceCache::Shutdown() {
    for (auto& entry : m_meshCache) {
        ShutdownMesh(entry.second);
    }
    m_meshCache.clear();

    for (auto& entry : m_textureCache) {
        ShutdownTexture(entry.second);
    }
    m_textureCache.clear();
}

const RenderResourceCache::CachedMesh* RenderResourceCache::GetOrCreateMesh(
    RHI::RHIDevice* device,
    const Resources::Mesh* mesh) {
    if (device == nullptr || mesh == nullptr || !mesh->IsValid()) {
        return nullptr;
    }

    auto it = m_meshCache.find(mesh);
    if (it != m_meshCache.end()) {
        return &it->second;
    }

    CachedMesh cachedMesh;
    if (!UploadMesh(device, mesh, cachedMesh)) {
        ShutdownMesh(cachedMesh);
        return nullptr;
    }

    const auto result = m_meshCache.emplace(mesh, cachedMesh);
    return &result.first->second;
}

const RenderResourceCache::CachedTexture* RenderResourceCache::GetOrCreateTexture(
    RHI::RHIDevice* device,
    const Resources::Texture* texture) {
    if (device == nullptr || texture == nullptr || !texture->IsValid()) {
        return nullptr;
    }

    auto it = m_textureCache.find(texture);
    if (it != m_textureCache.end()) {
        return &it->second;
    }

    CachedTexture cachedTexture;
    if (!UploadTexture(device, texture, cachedTexture)) {
        ShutdownTexture(cachedTexture);
        return nullptr;
    }

    const auto result = m_textureCache.emplace(texture, cachedTexture);
    return &result.first->second;
}

bool RenderResourceCache::UploadMesh(
    RHI::RHIDevice* device,
    const Resources::Mesh* mesh,
    CachedMesh& cachedMesh) {
    if (mesh->GetVertexData() == nullptr || mesh->GetVertexDataSize() == 0 || mesh->GetVertexStride() == 0) {
        return false;
    }

    RHI::BufferDesc vertexBufferDesc = {};
    vertexBufferDesc.size = static_cast<uint64_t>(mesh->GetVertexDataSize());
    vertexBufferDesc.stride = mesh->GetVertexStride();
    vertexBufferDesc.bufferType = static_cast<uint32_t>(RHI::BufferType::Vertex);
    cachedMesh.vertexBuffer = device->CreateBuffer(vertexBufferDesc);
    if (cachedMesh.vertexBuffer == nullptr) {
        return false;
    }

    cachedMesh.vertexBuffer->SetData(mesh->GetVertexData(), mesh->GetVertexDataSize());
    cachedMesh.vertexBuffer->SetStride(mesh->GetVertexStride());
    cachedMesh.vertexBuffer->SetBufferType(RHI::BufferType::Vertex);

    RHI::ResourceViewDesc vertexViewDesc = {};
    vertexViewDesc.dimension = RHI::ResourceViewDimension::Buffer;
    vertexViewDesc.structureByteStride = mesh->GetVertexStride();
    cachedMesh.vertexBufferView = device->CreateVertexBufferView(cachedMesh.vertexBuffer, vertexViewDesc);
    if (cachedMesh.vertexBufferView == nullptr) {
        return false;
    }

    cachedMesh.vertexCount = mesh->GetVertexCount();
    cachedMesh.vertexStride = mesh->GetVertexStride();

    if (mesh->GetIndexData() != nullptr && mesh->GetIndexDataSize() > 0 && mesh->GetIndexCount() > 0) {
        RHI::BufferDesc indexBufferDesc = {};
        indexBufferDesc.size = static_cast<uint64_t>(mesh->GetIndexDataSize());
        indexBufferDesc.stride = mesh->IsUse32BitIndex() ? sizeof(uint32_t) : sizeof(uint16_t);
        indexBufferDesc.bufferType = static_cast<uint32_t>(RHI::BufferType::Index);
        cachedMesh.indexBuffer = device->CreateBuffer(indexBufferDesc);
        if (cachedMesh.indexBuffer == nullptr) {
            return false;
        }

        cachedMesh.indexBuffer->SetData(mesh->GetIndexData(), mesh->GetIndexDataSize());
        cachedMesh.indexBuffer->SetStride(indexBufferDesc.stride);
        cachedMesh.indexBuffer->SetBufferType(RHI::BufferType::Index);

        RHI::ResourceViewDesc indexViewDesc = {};
        indexViewDesc.dimension = RHI::ResourceViewDimension::Buffer;
        indexViewDesc.format = static_cast<uint32_t>(
            mesh->IsUse32BitIndex() ? RHI::Format::R32_UInt : RHI::Format::R16_UInt);
        cachedMesh.indexBufferView = device->CreateIndexBufferView(cachedMesh.indexBuffer, indexViewDesc);
        if (cachedMesh.indexBufferView == nullptr) {
            return false;
        }

        cachedMesh.indexCount = mesh->GetIndexCount();
        cachedMesh.uses32BitIndices = mesh->IsUse32BitIndex();
    }

    return true;
}

bool RenderResourceCache::UploadTexture(
    RHI::RHIDevice* device,
    const Resources::Texture* texture,
    CachedTexture& cachedTexture) {
    if (texture->GetPixelData() == nullptr || texture->GetPixelDataSize() == 0) {
        return false;
    }

    const RHI::Format format = ToRHITextureFormat(texture->GetFormat());
    if (format == RHI::Format::Unknown) {
        return false;
    }

    RHI::TextureDesc textureDesc = {};
    textureDesc.width = texture->GetWidth();
    textureDesc.height = texture->GetHeight();
    textureDesc.depth = texture->GetDepth();
    textureDesc.mipLevels = texture->GetMipLevels();
    textureDesc.arraySize = 1;
    textureDesc.format = static_cast<uint32_t>(format);
    textureDesc.textureType = static_cast<uint32_t>(ToRHITextureType(texture->GetTextureType()));
    textureDesc.sampleCount = 1;
    textureDesc.sampleQuality = 0;
    textureDesc.flags = 0;

    const uint32_t rowPitch = texture->GetWidth() * 4;
    cachedTexture.texture = device->CreateTexture(
        textureDesc,
        texture->GetPixelData(),
        texture->GetPixelDataSize(),
        rowPitch);
    if (cachedTexture.texture == nullptr) {
        return false;
    }

    RHI::ResourceViewDesc textureViewDesc = {};
    textureViewDesc.format = static_cast<uint32_t>(format);
    textureViewDesc.dimension = RHI::ResourceViewDimension::Texture2D;
    textureViewDesc.mipLevel = 0;
    cachedTexture.shaderResourceView = device->CreateShaderResourceView(cachedTexture.texture, textureViewDesc);
    if (cachedTexture.shaderResourceView == nullptr) {
        return false;
    }

    cachedTexture.width = texture->GetWidth();
    cachedTexture.height = texture->GetHeight();
    return true;
}

} // namespace Rendering
} // namespace XCEngine