#pragma once

#include "RenderingIntegrationImageAssert.h"
#include "RenderingIntegrationMain.h"

#include <XCEngine/Components/CameraComponent.h>
#include <XCEngine/Components/GameObject.h>
#include <XCEngine/Core/Asset/IResource.h>
#include <XCEngine/Core/Math/Color.h>
#include <XCEngine/Core/Math/Vector3.h>
#include <XCEngine/Core/Math/Vector4.h>
#include <XCEngine/Rendering/Execution/SceneRenderer.h>
#include <XCEngine/Rendering/RenderContext.h>
#include <XCEngine/Rendering/RenderSurface.h>
#include <XCEngine/Resources/BuiltinResources.h>
#include <XCEngine/Resources/Material/Material.h>
#include <XCEngine/Resources/Mesh/Mesh.h>
#include <XCEngine/Resources/Shader/Shader.h>
#include <XCEngine/Resources/Texture/Texture.h>
#include <XCEngine/Resources/Volume/VolumeField.h>
#include <XCEngine/Resources/Volume/VolumeFieldLoader.h>
#include <XCEngine/RHI/RHITexture.h>
#include <XCEngine/Scene/Scene.h>

#include "../../RHI/integration/fixtures/RHIIntegrationFixture.h"

#include <filesystem>
#include <chrono>
#include <limits>
#include <memory>
#include <vector>

namespace VolumeIntegrationTestUtils {

using namespace XCEngine::Components;
using namespace XCEngine::Math;
using namespace XCEngine::Rendering;
using namespace XCEngine::Resources;
using namespace XCEngine::RHI;
using namespace XCEngine::RHI::Integration;

constexpr uint32_t kFrameWidth = 1280;
constexpr uint32_t kFrameHeight = 720;
constexpr const char* kCloudVolumeRelativePath = "Res/Volumes/cloud.nvdb";

struct StableFrameStats {
    int warmupFrames = 0;
    int measuredFrames = 0;
    double averageFrameTimeMs = 0.0;
    double minFrameTimeMs = 0.0;
    double maxFrameTimeMs = 0.0;
    double averageFps = 0.0;
};

inline Material* CreateVolumetricMaterial(
    const char* name,
    const char* path,
    const Vector4& tint = Vector4(1.0f, 1.0f, 1.0f, 1.0f),
    float densityScale = 0.2f,
    float stepSize = 1.0f,
    float maxSteps = 2000.0f,
    float ambientStrength = 0.005f,
    const Vector3& lightDirection = Vector3(0.5f, 0.8f, 0.3f),
    float lightSamples = 8.0f,
    float enableEntryHdda = 1.0f,
    float enableEmptySpaceSkipping = 1.0f,
    float enableEarlyTermination = 1.0f) {
    auto* material = new Material();
    IResource::ConstructParams params = {};
    params.name = name;
    params.path = path;
    params.guid = ResourceGUID::Generate(path);
    material->Initialize(params);
    material->SetShader(ResourceManager::Get().Load<Shader>(GetBuiltinVolumetricShaderPath()));
    material->SetRenderQueue(MaterialRenderQueue::Transparent);
    material->SetFloat4("_Tint", tint);
    material->SetFloat("_DensityScale", densityScale);
    material->SetFloat("_StepSize", stepSize);
    material->SetFloat("_MaxSteps", maxSteps);
    material->SetFloat("_AmbientStrength", ambientStrength);
    material->SetFloat4("_LightDirection", Vector4(lightDirection, 0.0f));
    material->SetFloat("_LightSamples", lightSamples);
    material->SetFloat("_EnableEntryHdda", enableEntryHdda);
    material->SetFloat("_EnableEmptySpaceSkipping", enableEmptySpaceSkipping);
    material->SetFloat("_EnableEarlyTermination", enableEarlyTermination);
    return material;
}

inline Material* CreateUnlitMaterial(
    const char* name,
    const char* path,
    const Vector4& baseColor) {
    static Texture* whiteTexture = []() -> Texture* {
        auto* texture = new Texture();
        IResource::ConstructParams params = {};
        params.name = "IntegrationWhiteTexture";
        params.path = "Tests/Rendering/Integration/white.texture";
        params.guid = ResourceGUID::Generate(params.path);
        texture->Initialize(params);

        const unsigned char pixel[4] = { 255, 255, 255, 255 };
        const bool created = texture->Create(
            1u,
            1u,
            1u,
            1u,
            XCEngine::Resources::TextureType::Texture2D,
            XCEngine::Resources::TextureFormat::RGBA8_UNORM,
            pixel,
            sizeof(pixel),
            1u);
        EXPECT_TRUE(created);
        return texture;
    }();

    auto* material = new Material();
    IResource::ConstructParams params = {};
    params.name = name;
    params.path = path;
    params.guid = ResourceGUID::Generate(path);
    material->Initialize(params);
    material->SetShader(ResourceManager::Get().Load<Shader>(GetBuiltinUnlitShaderPath()));
    material->SetRenderQueue(MaterialRenderQueue::Geometry);
    material->SetFloat4("_BaseColor", baseColor);
    material->SetTexture("_MainTex", ResourceHandle<Texture>(whiteTexture));
    return material;
}

inline VolumeField* LoadCloudVolumeField(const char* relativePath = kCloudVolumeRelativePath) {
    const std::filesystem::path volumePath =
        RenderingIntegrationTestUtils::ResolveRuntimePath(relativePath);
    EXPECT_TRUE(std::filesystem::exists(volumePath)) << volumePath.string();
    if (!std::filesystem::exists(volumePath)) {
        return nullptr;
    }

    VolumeFieldLoader loader;
    LoadResult result = loader.Load(volumePath.string().c_str());
    EXPECT_TRUE(result);
    EXPECT_NE(result.resource, nullptr);
    return result ? static_cast<VolumeField*>(result.resource) : nullptr;
}

inline ResourceHandle<Mesh> LoadBuiltinPrimitiveMesh(BuiltinPrimitiveType primitiveType) {
    return ResourceManager::Get().Load<Mesh>(GetBuiltinPrimitiveMeshPath(primitiveType));
}

class VolumeIntegrationSceneFixture : public RHIIntegrationFixture {
protected:
    virtual const char* GetSceneName() const = 0;
    virtual void BuildScene() = 0;
    virtual void ReleaseSceneResources() {}

    void SetUp() override {
        RHIIntegrationFixture::SetUp();

        mSceneRenderer = std::make_unique<SceneRenderer>();
        mScene = std::make_unique<Scene>(GetSceneName());

        BuildScene();

        TextureDesc depthDesc = {};
        depthDesc.width = kFrameWidth;
        depthDesc.height = kFrameHeight;
        depthDesc.depth = 1;
        depthDesc.mipLevels = 1;
        depthDesc.arraySize = 1;
        depthDesc.format = static_cast<uint32_t>(Format::D24_UNorm_S8_UInt);
        depthDesc.textureType = static_cast<uint32_t>(XCEngine::RHI::TextureType::Texture2D);
        depthDesc.sampleCount = 1;
        depthDesc.sampleQuality = 0;
        depthDesc.flags = 0;
        mDepthTexture = GetDevice()->CreateTexture(depthDesc);
        ASSERT_NE(mDepthTexture, nullptr);

        ResourceViewDesc depthViewDesc = {};
        depthViewDesc.format = static_cast<uint32_t>(Format::D24_UNorm_S8_UInt);
        depthViewDesc.dimension = ResourceViewDimension::Texture2D;
        depthViewDesc.mipLevel = 0;
        mDepthView = GetDevice()->CreateDepthStencilView(mDepthTexture, depthViewDesc);
        ASSERT_NE(mDepthView, nullptr);

        mBackBufferViews.resize(2, nullptr);
    }

    void TearDown() override {
        mSceneRenderer.reset();

        if (mDepthView != nullptr) {
            mDepthView->Shutdown();
            delete mDepthView;
            mDepthView = nullptr;
        }

        if (mDepthTexture != nullptr) {
            mDepthTexture->Shutdown();
            delete mDepthTexture;
            mDepthTexture = nullptr;
        }

        for (RHIResourceView*& backBufferView : mBackBufferViews) {
            if (backBufferView != nullptr) {
                backBufferView->Shutdown();
                delete backBufferView;
                backBufferView = nullptr;
            }
        }
        mBackBufferViews.clear();

        mScene.reset();
        ReleaseSceneResources();

        RHIIntegrationFixture::TearDown();
    }

    void RenderFrame() override {
        ASSERT_NE(mScene, nullptr);
        ASSERT_NE(mSceneRenderer, nullptr);

        RHICommandList* commandList = GetCommandList();
        ASSERT_NE(commandList, nullptr);

        commandList->Reset();

        RenderSurface surface(kFrameWidth, kFrameHeight);
        surface.SetColorAttachment(GetCurrentBackBufferView());
        surface.SetDepthAttachment(mDepthView);

        RenderContext renderContext = {};
        renderContext.device = GetDevice();
        renderContext.commandList = commandList;
        renderContext.commandQueue = GetCommandQueue();
        renderContext.backendType = GetBackendType();

        ASSERT_TRUE(mSceneRenderer->Render(*mScene, nullptr, renderContext, surface));

        commandList->Close();
        void* commandLists[] = { commandList };
        GetCommandQueue()->ExecuteCommandLists(1, commandLists);
    }

    void RenderAndCompare(const char* screenshotFilename, float comparisonThreshold = 0.0f) {
        RHICommandQueue* commandQueue = GetCommandQueue();
        RHISwapChain* swapChain = GetSwapChain();
        ASSERT_NE(swapChain, nullptr);

        constexpr int kTargetFrameCount = 30;
        for (int frameCount = 0; frameCount <= kTargetFrameCount; ++frameCount) {
            if (frameCount > 0) {
                commandQueue->WaitForPreviousFrame();
            }

            BeginRender();
            RenderFrame();

            if (frameCount >= kTargetFrameCount) {
                commandQueue->WaitForIdle();
                ASSERT_TRUE(TakeScreenshot(screenshotFilename));
                ASSERT_TRUE(CompareWithGoldenTemplate(
                    screenshotFilename,
                    "GT.ppm",
                    comparisonThreshold));
                break;
            }

            swapChain->Present(0, 0);
        }
    }

    StableFrameStats RenderAndMeasureStableFrames(
        int warmupFrameCount,
        int measuredFrameCount,
        const char* screenshotFilename = nullptr) {
        StableFrameStats stats = {};
        stats.warmupFrames = warmupFrameCount;
        stats.measuredFrames = measuredFrameCount;

        RHICommandQueue* commandQueue = GetCommandQueue();
        RHISwapChain* swapChain = GetSwapChain();
        if (commandQueue == nullptr || swapChain == nullptr) {
            ADD_FAILURE() << "RenderAndMeasureStableFrames requires a valid command queue and swap chain";
            return stats;
        }

        auto renderPresentedFrame = [&]() {
            BeginRender();
            RenderFrame();
            swapChain->Present(0, 0);
        };

        for (int frameIndex = 0; frameIndex < warmupFrameCount; ++frameIndex) {
            if (frameIndex > 0) {
                commandQueue->WaitForPreviousFrame();
            }
            renderPresentedFrame();
        }

        if (warmupFrameCount > 0) {
            commandQueue->WaitForPreviousFrame();
        }

        using Clock = std::chrono::steady_clock;
        double totalFrameTimeMs = 0.0;
        double minFrameTimeMs = std::numeric_limits<double>::max();
        double maxFrameTimeMs = 0.0;
        for (int frameIndex = 0; frameIndex < measuredFrameCount; ++frameIndex) {
            const Clock::time_point frameStart = Clock::now();
            renderPresentedFrame();
            commandQueue->WaitForPreviousFrame();
            const double frameTimeMs = std::chrono::duration<double, std::milli>(
                Clock::now() - frameStart).count();
            totalFrameTimeMs += frameTimeMs;
            minFrameTimeMs = (std::min)(minFrameTimeMs, frameTimeMs);
            maxFrameTimeMs = (std::max)(maxFrameTimeMs, frameTimeMs);
        }

        if (measuredFrameCount > 0) {
            stats.averageFrameTimeMs = totalFrameTimeMs / static_cast<double>(measuredFrameCount);
            stats.minFrameTimeMs = minFrameTimeMs;
            stats.maxFrameTimeMs = maxFrameTimeMs;
            stats.averageFps = stats.averageFrameTimeMs > 0.0
                ? 1000.0 / stats.averageFrameTimeMs
                : 0.0;
        }

        if (screenshotFilename != nullptr && screenshotFilename[0] != '\0') {
            BeginRender();
            RenderFrame();
            commandQueue->WaitForIdle();
            EXPECT_TRUE(TakeScreenshot(screenshotFilename));
        }

        return stats;
    }

    RHIResourceView* GetCurrentBackBufferView() {
        const int backBufferIndex = GetCurrentBackBufferIndex();
        if (backBufferIndex < 0) {
            return nullptr;
        }

        if (static_cast<size_t>(backBufferIndex) >= mBackBufferViews.size()) {
            mBackBufferViews.resize(static_cast<size_t>(backBufferIndex) + 1, nullptr);
        }

        if (mBackBufferViews[backBufferIndex] == nullptr) {
            ResourceViewDesc viewDesc = {};
            viewDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
            viewDesc.dimension = ResourceViewDimension::Texture2D;
            viewDesc.mipLevel = 0;
            mBackBufferViews[backBufferIndex] =
                GetDevice()->CreateRenderTargetView(GetCurrentBackBuffer(), viewDesc);
        }

        return mBackBufferViews[backBufferIndex];
    }

    std::unique_ptr<Scene> mScene;
    std::unique_ptr<SceneRenderer> mSceneRenderer;

private:
    std::vector<RHIResourceView*> mBackBufferViews;
    RHITexture* mDepthTexture = nullptr;
    RHIResourceView* mDepthView = nullptr;
};

} // namespace VolumeIntegrationTestUtils