#include <gtest/gtest.h>

#include <XCEngine/Components/GaussianSplatRendererComponent.h>
#include <XCEngine/Components/GameObject.h>
#include <XCEngine/Core/Asset/IResource.h>
#include <XCEngine/Resources/GaussianSplat/GaussianSplat.h>
#include <XCEngine/RHI/RHIResourceView.h>
#include <XCEngine/RHI/RHITexture.h>

#include "Rendering/Passes/Internal/BuiltinGaussianSplatPassResources.h"

#include <cstring>
#include <memory>
#include <string>

using namespace XCEngine::Components;
using namespace XCEngine::Rendering;
using namespace XCEngine::Rendering::Passes::Internal;
using namespace XCEngine::Resources;

namespace {

struct MockGaussianSplatResourceState {
    int createBufferCalls = 0;
    int bufferShutdownCalls = 0;
    int bufferDestroyCalls = 0;
    int createBufferShaderViewCalls = 0;
    int createBufferUavCalls = 0;
    int createTextureCalls = 0;
    int textureShutdownCalls = 0;
    int textureDestroyCalls = 0;
    int createTextureShaderViewCalls = 0;
    int createRenderTargetViewCalls = 0;
    int resourceViewShutdownCalls = 0;
    int resourceViewDestroyCalls = 0;
};

class MockGaussianSplatBuffer final : public XCEngine::RHI::RHIBuffer {
public:
    MockGaussianSplatBuffer(
        std::shared_ptr<MockGaussianSplatResourceState> state,
        const XCEngine::RHI::BufferDesc& desc)
        : m_state(std::move(state))
        , m_size(desc.size)
        , m_stride(desc.stride)
        , m_type(static_cast<XCEngine::RHI::BufferType>(desc.bufferType)) {
    }

    ~MockGaussianSplatBuffer() override {
        ++m_state->bufferDestroyCalls;
    }

    void* Map() override { return nullptr; }
    void Unmap() override {}
    void SetData(const void*, size_t, size_t) override {}
    uint64_t GetSize() const override { return m_size; }
    XCEngine::RHI::BufferType GetBufferType() const override { return m_type; }
    void SetBufferType(XCEngine::RHI::BufferType type) override { m_type = type; }
    uint32_t GetStride() const override { return m_stride; }
    void SetStride(uint32_t stride) override { m_stride = stride; }
    void* GetNativeHandle() override { return nullptr; }
    XCEngine::RHI::ResourceStates GetState() const override { return m_resourceState; }
    void SetState(XCEngine::RHI::ResourceStates state) override { m_resourceState = state; }
    const std::string& GetName() const override { return m_name; }
    void SetName(const std::string& name) override { m_name = name; }

    void Shutdown() override {
        ++m_state->bufferShutdownCalls;
    }

private:
    std::shared_ptr<MockGaussianSplatResourceState> m_state;
    uint64_t m_size = 0u;
    uint32_t m_stride = 0u;
    XCEngine::RHI::BufferType m_type = XCEngine::RHI::BufferType::Storage;
    XCEngine::RHI::ResourceStates m_resourceState = XCEngine::RHI::ResourceStates::Common;
    std::string m_name;
};

class MockGaussianSplatTexture final : public XCEngine::RHI::RHITexture {
public:
    MockGaussianSplatTexture(
        std::shared_ptr<MockGaussianSplatResourceState> state,
        const XCEngine::RHI::TextureDesc& desc)
        : m_state(std::move(state))
        , m_width(desc.width)
        , m_height(desc.height)
        , m_format(static_cast<XCEngine::RHI::Format>(desc.format))
        , m_type(static_cast<XCEngine::RHI::TextureType>(desc.textureType)) {
    }

    ~MockGaussianSplatTexture() override {
        ++m_state->textureDestroyCalls;
    }

    uint32_t GetWidth() const override { return m_width; }
    uint32_t GetHeight() const override { return m_height; }
    uint32_t GetDepth() const override { return 1u; }
    uint32_t GetMipLevels() const override { return 1u; }
    XCEngine::RHI::Format GetFormat() const override { return m_format; }
    XCEngine::RHI::TextureType GetTextureType() const override { return m_type; }
    XCEngine::RHI::ResourceStates GetState() const override { return m_resourceState; }
    void SetState(XCEngine::RHI::ResourceStates state) override { m_resourceState = state; }
    void* GetNativeHandle() override { return nullptr; }
    const std::string& GetName() const override { return m_name; }
    void SetName(const std::string& name) override { m_name = name; }

    void Shutdown() override {
        ++m_state->textureShutdownCalls;
    }

private:
    std::shared_ptr<MockGaussianSplatResourceState> m_state;
    uint32_t m_width = 0u;
    uint32_t m_height = 0u;
    XCEngine::RHI::Format m_format = XCEngine::RHI::Format::Unknown;
    XCEngine::RHI::TextureType m_type = XCEngine::RHI::TextureType::Texture2D;
    XCEngine::RHI::ResourceStates m_resourceState = XCEngine::RHI::ResourceStates::Common;
    std::string m_name;
};

class MockGaussianSplatResourceView final : public XCEngine::RHI::RHIResourceView {
public:
    MockGaussianSplatResourceView(
        std::shared_ptr<MockGaussianSplatResourceState> state,
        XCEngine::RHI::ResourceViewType type,
        XCEngine::RHI::Format format,
        XCEngine::RHI::ResourceViewDimension dimension)
        : m_state(std::move(state))
        , m_type(type)
        , m_format(format)
        , m_dimension(dimension) {
    }

    ~MockGaussianSplatResourceView() override {
        ++m_state->resourceViewDestroyCalls;
    }

    void Shutdown() override {
        ++m_state->resourceViewShutdownCalls;
    }

    void* GetNativeHandle() override { return nullptr; }
    bool IsValid() const override { return true; }
    XCEngine::RHI::ResourceViewType GetViewType() const override { return m_type; }
    XCEngine::RHI::ResourceViewDimension GetDimension() const override { return m_dimension; }
    XCEngine::RHI::Format GetFormat() const override { return m_format; }

private:
    std::shared_ptr<MockGaussianSplatResourceState> m_state;
    XCEngine::RHI::ResourceViewType m_type = XCEngine::RHI::ResourceViewType::ShaderResource;
    XCEngine::RHI::Format m_format = XCEngine::RHI::Format::Unknown;
    XCEngine::RHI::ResourceViewDimension m_dimension = XCEngine::RHI::ResourceViewDimension::Unknown;
};

class MockGaussianSplatDevice final : public XCEngine::RHI::RHIDevice {
public:
    explicit MockGaussianSplatDevice(std::shared_ptr<MockGaussianSplatResourceState> state)
        : m_state(std::move(state)) {
    }

    bool Initialize(const XCEngine::RHI::RHIDeviceDesc&) override { return true; }
    void Shutdown() override {}

    XCEngine::RHI::RHIBuffer* CreateBuffer(const XCEngine::RHI::BufferDesc& desc) override {
        ++m_state->createBufferCalls;
        return new MockGaussianSplatBuffer(m_state, desc);
    }

    XCEngine::RHI::RHITexture* CreateTexture(const XCEngine::RHI::TextureDesc& desc) override {
        ++m_state->createTextureCalls;
        return new MockGaussianSplatTexture(m_state, desc);
    }

    XCEngine::RHI::RHITexture* CreateTexture(
        const XCEngine::RHI::TextureDesc& desc,
        const void*,
        size_t,
        uint32_t) override {
        return CreateTexture(desc);
    }

    XCEngine::RHI::RHISwapChain* CreateSwapChain(
        const XCEngine::RHI::SwapChainDesc&,
        XCEngine::RHI::RHICommandQueue*) override { return nullptr; }
    XCEngine::RHI::RHICommandList* CreateCommandList(const XCEngine::RHI::CommandListDesc&) override { return nullptr; }
    XCEngine::RHI::RHICommandQueue* CreateCommandQueue(const XCEngine::RHI::CommandQueueDesc&) override { return nullptr; }
    XCEngine::RHI::RHIShader* CreateShader(const XCEngine::RHI::ShaderCompileDesc&) override { return nullptr; }
    XCEngine::RHI::RHIPipelineState* CreatePipelineState(const XCEngine::RHI::GraphicsPipelineDesc&) override { return nullptr; }
    XCEngine::RHI::RHIPipelineLayout* CreatePipelineLayout(const XCEngine::RHI::RHIPipelineLayoutDesc&) override { return nullptr; }
    XCEngine::RHI::RHIFence* CreateFence(const XCEngine::RHI::FenceDesc&) override { return nullptr; }
    XCEngine::RHI::RHISampler* CreateSampler(const XCEngine::RHI::SamplerDesc&) override { return nullptr; }
    XCEngine::RHI::RHIRenderPass* CreateRenderPass(
        uint32_t,
        const XCEngine::RHI::AttachmentDesc*,
        const XCEngine::RHI::AttachmentDesc*) override { return nullptr; }
    XCEngine::RHI::RHIFramebuffer* CreateFramebuffer(
        XCEngine::RHI::RHIRenderPass*,
        uint32_t,
        uint32_t,
        uint32_t,
        XCEngine::RHI::RHIResourceView**,
        XCEngine::RHI::RHIResourceView*) override { return nullptr; }
    XCEngine::RHI::RHIDescriptorPool* CreateDescriptorPool(const XCEngine::RHI::DescriptorPoolDesc&) override { return nullptr; }
    XCEngine::RHI::RHIDescriptorSet* CreateDescriptorSet(
        XCEngine::RHI::RHIDescriptorPool*,
        const XCEngine::RHI::DescriptorSetLayoutDesc&) override { return nullptr; }
    XCEngine::RHI::RHIResourceView* CreateVertexBufferView(
        XCEngine::RHI::RHIBuffer*,
        const XCEngine::RHI::ResourceViewDesc&) override { return nullptr; }
    XCEngine::RHI::RHIResourceView* CreateIndexBufferView(
        XCEngine::RHI::RHIBuffer*,
        const XCEngine::RHI::ResourceViewDesc&) override { return nullptr; }

    XCEngine::RHI::RHIResourceView* CreateRenderTargetView(
        XCEngine::RHI::RHITexture*,
        const XCEngine::RHI::ResourceViewDesc& desc) override {
        ++m_state->createRenderTargetViewCalls;
        return new MockGaussianSplatResourceView(
            m_state,
            XCEngine::RHI::ResourceViewType::RenderTarget,
            static_cast<XCEngine::RHI::Format>(desc.format),
            desc.dimension);
    }

    XCEngine::RHI::RHIResourceView* CreateDepthStencilView(
        XCEngine::RHI::RHITexture*,
        const XCEngine::RHI::ResourceViewDesc&) override { return nullptr; }

    XCEngine::RHI::RHIResourceView* CreateShaderResourceView(
        XCEngine::RHI::RHIBuffer*,
        const XCEngine::RHI::ResourceViewDesc& desc) override {
        ++m_state->createBufferShaderViewCalls;
        return new MockGaussianSplatResourceView(
            m_state,
            XCEngine::RHI::ResourceViewType::ShaderResource,
            static_cast<XCEngine::RHI::Format>(desc.format),
            desc.dimension);
    }

    XCEngine::RHI::RHIResourceView* CreateShaderResourceView(
        XCEngine::RHI::RHITexture*,
        const XCEngine::RHI::ResourceViewDesc& desc) override {
        ++m_state->createTextureShaderViewCalls;
        return new MockGaussianSplatResourceView(
            m_state,
            XCEngine::RHI::ResourceViewType::ShaderResource,
            static_cast<XCEngine::RHI::Format>(desc.format),
            desc.dimension);
    }

    XCEngine::RHI::RHIResourceView* CreateUnorderedAccessView(
        XCEngine::RHI::RHIBuffer*,
        const XCEngine::RHI::ResourceViewDesc& desc) override {
        ++m_state->createBufferUavCalls;
        return new MockGaussianSplatResourceView(
            m_state,
            XCEngine::RHI::ResourceViewType::UnorderedAccess,
            static_cast<XCEngine::RHI::Format>(desc.format),
            desc.dimension);
    }

    XCEngine::RHI::RHIResourceView* CreateUnorderedAccessView(
        XCEngine::RHI::RHITexture*,
        const XCEngine::RHI::ResourceViewDesc&) override { return nullptr; }

    const XCEngine::RHI::RHICapabilities& GetCapabilities() const override { return m_capabilities; }
    const XCEngine::RHI::RHIDeviceInfo& GetDeviceInfo() const override { return m_info; }
    void* GetNativeDevice() override { return nullptr; }

private:
    std::shared_ptr<MockGaussianSplatResourceState> m_state;
    XCEngine::RHI::RHICapabilities m_capabilities = {};
    XCEngine::RHI::RHIDeviceInfo m_info = {};
};

GaussianSplat* CreateTestGaussianSplat(const char* path, XCEngine::Core::uint32 splatCount) {
    auto* gaussianSplat = new GaussianSplat();
    IResource::ConstructParams params = {};
    params.name = "TestGaussianSplat";
    params.path = path;
    params.guid = ResourceGUID::Generate(path);
    gaussianSplat->Initialize(params);

    GaussianSplatMetadata metadata = {};
    metadata.splatCount = splatCount;

    XCEngine::Containers::Array<GaussianSplatSection> sections;
    sections.Resize(1);
    sections[0].type = GaussianSplatSectionType::Positions;
    sections[0].format = GaussianSplatSectionFormat::VectorFloat32;
    sections[0].dataOffset = 0u;
    sections[0].dataSize = static_cast<XCEngine::Core::uint64>(splatCount) * sizeof(GaussianSplatPositionRecord);
    sections[0].elementCount = splatCount;
    sections[0].elementStride = sizeof(GaussianSplatPositionRecord);

    XCEngine::Containers::Array<XCEngine::Core::uint8> payload;
    payload.Resize(static_cast<size_t>(sections[0].dataSize));
    for (XCEngine::Core::uint32 index = 0; index < splatCount; ++index) {
        const GaussianSplatPositionRecord positionRecord = { XCEngine::Math::Vector3(static_cast<float>(index), 0.0f, 0.0f) };
        std::memcpy(
            payload.Data() + (static_cast<size_t>(index) * sizeof(GaussianSplatPositionRecord)),
            &positionRecord,
            sizeof(positionRecord));
    }

    EXPECT_TRUE(gaussianSplat->CreateOwned(metadata, std::move(sections), std::move(payload)));
    return gaussianSplat;
}

VisibleGaussianSplatItem BuildVisibleGaussianSplatItem(
    GameObject& gameObject,
    GaussianSplatRendererComponent& renderer,
    GaussianSplat& gaussianSplat) {
    VisibleGaussianSplatItem item = {};
    item.gameObject = &gameObject;
    item.gaussianSplatRenderer = &renderer;
    item.gaussianSplat = &gaussianSplat;
    return item;
}

TEST(BuiltinGaussianSplatPassResources_Test, EnsureWorkingSetAllocatesAndReusesStructuredBuffersPerRenderer) {
    auto state = std::make_shared<MockGaussianSplatResourceState>();
    MockGaussianSplatDevice device(state);
    BuiltinGaussianSplatPassResources resources;

    GameObject gameObject("GaussianSplatObject");
    auto* renderer = gameObject.AddComponent<GaussianSplatRendererComponent>();
    std::unique_ptr<GaussianSplat> gaussianSplat(CreateTestGaussianSplat("GaussianSplats/room.xcgsplat", 8u));
    VisibleGaussianSplatItem item = BuildVisibleGaussianSplatItem(gameObject, *renderer, *gaussianSplat);

    BuiltinGaussianSplatPassResources::WorkingSet* workingSet = nullptr;
    ASSERT_TRUE(resources.EnsureWorkingSet(&device, item, workingSet));
    ASSERT_NE(workingSet, nullptr);
    EXPECT_EQ(resources.GetWorkingSetCount(), 1u);
    EXPECT_EQ(workingSet->renderer, renderer);
    EXPECT_EQ(workingSet->splatCapacity, 8u);
    EXPECT_EQ(workingSet->sortDistances.elementStride, sizeof(float));
    EXPECT_EQ(workingSet->orderIndices.elementStride, sizeof(XCEngine::Core::uint32));
    EXPECT_EQ(workingSet->viewData.elementStride, sizeof(GaussianSplatViewData));
    EXPECT_EQ(workingSet->sortDistances.shaderResourceView->GetDimension(), XCEngine::RHI::ResourceViewDimension::StructuredBuffer);
    EXPECT_EQ(workingSet->sortDistances.unorderedAccessView->GetViewType(), XCEngine::RHI::ResourceViewType::UnorderedAccess);
    EXPECT_EQ(state->createBufferCalls, 3);
    EXPECT_EQ(state->createBufferShaderViewCalls, 3);
    EXPECT_EQ(state->createBufferUavCalls, 3);

    BuiltinGaussianSplatPassResources::WorkingSet* reusedWorkingSet = nullptr;
    ASSERT_TRUE(resources.EnsureWorkingSet(&device, item, reusedWorkingSet));
    EXPECT_EQ(reusedWorkingSet, workingSet);
    EXPECT_EQ(state->createBufferCalls, 3);
    EXPECT_EQ(state->createBufferShaderViewCalls, 3);
    EXPECT_EQ(state->createBufferUavCalls, 3);
}

TEST(BuiltinGaussianSplatPassResources_Test, EnsureWorkingSetKeepsPerRendererIsolationAndRecreatesOnCapacityGrowth) {
    auto state = std::make_shared<MockGaussianSplatResourceState>();
    MockGaussianSplatDevice device(state);
    BuiltinGaussianSplatPassResources resources;

    GameObject firstObject("FirstGaussianSplat");
    auto* firstRenderer = firstObject.AddComponent<GaussianSplatRendererComponent>();
    std::unique_ptr<GaussianSplat> firstGaussianSplat(CreateTestGaussianSplat("GaussianSplats/first.xcgsplat", 4u));
    VisibleGaussianSplatItem firstItem = BuildVisibleGaussianSplatItem(firstObject, *firstRenderer, *firstGaussianSplat);

    GameObject secondObject("SecondGaussianSplat");
    auto* secondRenderer = secondObject.AddComponent<GaussianSplatRendererComponent>();
    std::unique_ptr<GaussianSplat> secondGaussianSplat(CreateTestGaussianSplat("GaussianSplats/second.xcgsplat", 4u));
    VisibleGaussianSplatItem secondItem = BuildVisibleGaussianSplatItem(secondObject, *secondRenderer, *secondGaussianSplat);

    BuiltinGaussianSplatPassResources::WorkingSet* firstWorkingSet = nullptr;
    BuiltinGaussianSplatPassResources::WorkingSet* secondWorkingSet = nullptr;
    ASSERT_TRUE(resources.EnsureWorkingSet(&device, firstItem, firstWorkingSet));
    ASSERT_TRUE(resources.EnsureWorkingSet(&device, secondItem, secondWorkingSet));
    ASSERT_NE(firstWorkingSet, nullptr);
    ASSERT_NE(secondWorkingSet, nullptr);
    EXPECT_NE(firstWorkingSet, secondWorkingSet);
    EXPECT_EQ(resources.GetWorkingSetCount(), 2u);
    EXPECT_EQ(state->createBufferCalls, 6);

    std::unique_ptr<GaussianSplat> grownGaussianSplat(CreateTestGaussianSplat("GaussianSplats/first_grown.xcgsplat", 12u));
    firstItem.gaussianSplat = grownGaussianSplat.get();

    BuiltinGaussianSplatPassResources::WorkingSet* grownWorkingSet = nullptr;
    ASSERT_TRUE(resources.EnsureWorkingSet(&device, firstItem, grownWorkingSet));
    ASSERT_NE(grownWorkingSet, nullptr);
    EXPECT_EQ(grownWorkingSet, resources.FindWorkingSet(firstRenderer));
    EXPECT_EQ(grownWorkingSet->splatCapacity, 12u);
    EXPECT_EQ(resources.GetWorkingSetCount(), 2u);
    EXPECT_EQ(state->createBufferCalls, 9);
    EXPECT_GE(state->bufferShutdownCalls, 3);
    EXPECT_GE(state->bufferDestroyCalls, 3);
}

TEST(BuiltinGaussianSplatPassResources_Test, EnsureAccumulationSurfaceReusesCompatibleTargetAndRecreatesOnResize) {
    auto state = std::make_shared<MockGaussianSplatResourceState>();
    MockGaussianSplatDevice device(state);
    BuiltinGaussianSplatPassResources resources;

    BuiltinGaussianSplatPassResources::AccumulationSurface* surface = nullptr;
    ASSERT_TRUE(resources.EnsureAccumulationSurface(
        &device,
        640u,
        360u,
        XCEngine::RHI::Format::R16G16B16A16_Float,
        surface));
    ASSERT_NE(surface, nullptr);
    EXPECT_EQ(surface->width, 640u);
    EXPECT_EQ(surface->height, 360u);
    EXPECT_EQ(surface->format, XCEngine::RHI::Format::R16G16B16A16_Float);
    EXPECT_EQ(state->createTextureCalls, 1);
    EXPECT_EQ(state->createRenderTargetViewCalls, 1);
    EXPECT_EQ(state->createTextureShaderViewCalls, 1);

    BuiltinGaussianSplatPassResources::AccumulationSurface* reusedSurface = nullptr;
    ASSERT_TRUE(resources.EnsureAccumulationSurface(
        &device,
        640u,
        360u,
        XCEngine::RHI::Format::R16G16B16A16_Float,
        reusedSurface));
    EXPECT_EQ(reusedSurface, surface);
    EXPECT_EQ(state->createTextureCalls, 1);

    BuiltinGaussianSplatPassResources::AccumulationSurface* resizedSurface = nullptr;
    ASSERT_TRUE(resources.EnsureAccumulationSurface(
        &device,
        800u,
        600u,
        XCEngine::RHI::Format::R16G16B16A16_Float,
        resizedSurface));
    ASSERT_NE(resizedSurface, nullptr);
    EXPECT_EQ(resizedSurface->width, 800u);
    EXPECT_EQ(resizedSurface->height, 600u);
    EXPECT_EQ(state->createTextureCalls, 2);
    EXPECT_GE(state->textureShutdownCalls, 1);
    EXPECT_GE(state->textureDestroyCalls, 1);
    EXPECT_GE(state->resourceViewShutdownCalls, 2);
    EXPECT_GE(state->resourceViewDestroyCalls, 2);
}

} // namespace