#define NOMINMAX
#include <windows.h>

#include <gtest/gtest.h>

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

#include <XCEngine/Components/CameraComponent.h>
#include <XCEngine/Components/GameObject.h>
#include <XCEngine/Components/LightComponent.h>
#include <XCEngine/Components/MeshFilterComponent.h>
#include <XCEngine/Components/MeshRendererComponent.h>
#include <XCEngine/Core/Asset/IResource.h>
#include <XCEngine/Core/Math/Color.h>
#include <XCEngine/Core/Math/Quaternion.h>
#include <XCEngine/Core/Math/Vector2.h>
#include <XCEngine/Core/Math/Vector3.h>
#include <XCEngine/Rendering/RenderContext.h>
#include <XCEngine/Rendering/RenderSurface.h>
#include <XCEngine/Rendering/SceneRenderer.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/RHI/RHITexture.h>
#include <XCEngine/Scene/Scene.h>

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

#include <memory>
#include <vector>

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;

namespace {

constexpr const char* kD3D12Screenshot = "multi_light_scene_d3d12.ppm";
constexpr const char* kOpenGLScreenshot = "multi_light_scene_opengl.ppm";
constexpr const char* kVulkanScreenshot = "multi_light_scene_vulkan.ppm";
constexpr uint32_t kFrameWidth = 1280;
constexpr uint32_t kFrameHeight = 720;

void ExpectMultiLightKeyPixels(const RenderingIntegrationTestUtils::PpmImage& image) {
    using namespace RenderingIntegrationTestUtils;

    ExpectPixelLuminanceAtMost(image, 320, 320, 24, "background remains dark");
    ExpectPixelLuminanceAtLeast(image, 625, 285, 540, "left point-light highlight");
    ExpectPixelChannelDominates(image, 625, 285, 0, 80, "left point-light highlight stays warm");
    ExpectPixelLuminanceAtLeast(image, 640, 315, 470, "center cube receives main light");
    ExpectPixelChannelDominates(image, 640, 315, 0, 40, "center cube keeps warm directional tint");
    ExpectPixelLuminanceAtLeast(image, 758, 311, 300, "right cube receives cool fill");
    ExpectPixelChannelDominates(image, 758, 311, 2, 20, "right cube keeps cool fill tint");
}

void AppendQuadFace(
    std::vector<StaticMeshVertex>& vertices,
    std::vector<uint32_t>& indices,
    const Vector3& v0,
    const Vector3& v1,
    const Vector3& v2,
    const Vector3& v3,
    const Vector3& normal) {
    const uint32_t baseIndex = static_cast<uint32_t>(vertices.size());

    StaticMeshVertex vertex = {};
    vertex.normal = normal;

    vertex.position = v0;
    vertex.uv0 = Vector2(0.0f, 1.0f);
    vertices.push_back(vertex);

    vertex.position = v1;
    vertex.uv0 = Vector2(1.0f, 1.0f);
    vertices.push_back(vertex);

    vertex.position = v2;
    vertex.uv0 = Vector2(0.0f, 0.0f);
    vertices.push_back(vertex);

    vertex.position = v3;
    vertex.uv0 = Vector2(1.0f, 0.0f);
    vertices.push_back(vertex);

    indices.push_back(baseIndex + 0);
    indices.push_back(baseIndex + 2);
    indices.push_back(baseIndex + 1);
    indices.push_back(baseIndex + 1);
    indices.push_back(baseIndex + 2);
    indices.push_back(baseIndex + 3);
}

Mesh* CreateCubeMesh() {
    auto* mesh = new Mesh();
    IResource::ConstructParams params = {};
    params.name = "MultiLightCubeMesh";
    params.path = "Tests/Rendering/MultiLightCube.mesh";
    params.guid = ResourceGUID::Generate(params.path);
    mesh->Initialize(params);

    std::vector<StaticMeshVertex> vertices;
    std::vector<uint32_t> indices;
    vertices.reserve(24);
    indices.reserve(36);

    constexpr float half = 0.5f;
    AppendQuadFace(
        vertices, indices,
        Vector3(-half, -half, half),
        Vector3(half, -half, half),
        Vector3(-half, half, half),
        Vector3(half, half, half),
        Vector3::Forward());
    AppendQuadFace(
        vertices, indices,
        Vector3(half, -half, -half),
        Vector3(-half, -half, -half),
        Vector3(half, half, -half),
        Vector3(-half, half, -half),
        Vector3::Back());
    AppendQuadFace(
        vertices, indices,
        Vector3(-half, -half, -half),
        Vector3(-half, -half, half),
        Vector3(-half, half, -half),
        Vector3(-half, half, half),
        Vector3::Left());
    AppendQuadFace(
        vertices, indices,
        Vector3(half, -half, half),
        Vector3(half, -half, -half),
        Vector3(half, half, half),
        Vector3(half, half, -half),
        Vector3::Right());
    AppendQuadFace(
        vertices, indices,
        Vector3(-half, half, half),
        Vector3(half, half, half),
        Vector3(-half, half, -half),
        Vector3(half, half, -half),
        Vector3::Up());
    AppendQuadFace(
        vertices, indices,
        Vector3(-half, -half, -half),
        Vector3(half, -half, -half),
        Vector3(-half, -half, half),
        Vector3(half, -half, half),
        Vector3::Down());

    mesh->SetVertexData(
        vertices.data(),
        vertices.size() * sizeof(StaticMeshVertex),
        static_cast<uint32_t>(vertices.size()),
        sizeof(StaticMeshVertex),
        VertexAttribute::Position | VertexAttribute::Normal | VertexAttribute::UV0);
    mesh->SetIndexData(
        indices.data(),
        indices.size() * sizeof(uint32_t),
        static_cast<uint32_t>(indices.size()),
        true);

    const Bounds bounds(Vector3::Zero(), Vector3::One());
    mesh->SetBounds(bounds);

    MeshSection section = {};
    section.baseVertex = 0;
    section.vertexCount = static_cast<uint32_t>(vertices.size());
    section.startIndex = 0;
    section.indexCount = static_cast<uint32_t>(indices.size());
    section.materialID = 0;
    section.bounds = bounds;
    mesh->AddSection(section);
    return mesh;
}

Material* CreateForwardLitMaterial(
    const char* name,
    const char* path,
    const Vector4& baseColor) {
    auto* material = new Material();
    IResource::ConstructParams params = {};
    params.name = name;
    params.path = path;
    params.guid = ResourceGUID::Generate(params.path);
    material->Initialize(params);
    material->SetShader(ResourceManager::Get().Load<Shader>(GetBuiltinForwardLitShaderPath()));
    material->SetShaderPass("ForwardLit");
    material->SetFloat4("_BaseColor", baseColor);
    return material;
}

const char* GetScreenshotFilename(RHIType backendType) {
    switch (backendType) {
        case RHIType::D3D12:
            return kD3D12Screenshot;
        case RHIType::Vulkan:
            return kVulkanScreenshot;
        case RHIType::OpenGL:
        default:
            return kOpenGLScreenshot;
    }
}

int GetComparisonThreshold(RHIType backendType) {
    return backendType == RHIType::D3D12 ? 10 : 10;
}

class MultiLightSceneTest : public RHIIntegrationFixture {
protected:
    void SetUp() override;
    void TearDown() override;
    void RenderFrame() override;

private:
    void BuildScene();
    RHIResourceView* GetCurrentBackBufferView();

    std::unique_ptr<Scene> mScene;
    std::unique_ptr<SceneRenderer> mSceneRenderer;
    std::vector<RHIResourceView*> mBackBufferViews;
    RHITexture* mDepthTexture = nullptr;
    RHIResourceView* mDepthView = nullptr;
    Mesh* mCubeMesh = nullptr;
    Material* mGroundMaterial = nullptr;
    Material* mLeftMaterial = nullptr;
    Material* mCenterMaterial = nullptr;
    Material* mRightMaterial = nullptr;
};

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

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

    mCubeMesh = CreateCubeMesh();
    ASSERT_NE(mCubeMesh, nullptr);

    mGroundMaterial = CreateForwardLitMaterial(
        "MultiLightGround",
        "Tests/Rendering/MultiLightGround.material",
        Vector4(0.84f, 0.84f, 0.86f, 1.0f));
    mLeftMaterial = CreateForwardLitMaterial(
        "MultiLightLeftCube",
        "Tests/Rendering/MultiLightLeftCube.material",
        Vector4(0.95f, 0.92f, 0.90f, 1.0f));
    mCenterMaterial = CreateForwardLitMaterial(
        "MultiLightCenterCube",
        "Tests/Rendering/MultiLightCenterCube.material",
        Vector4(0.92f, 0.92f, 0.95f, 1.0f));
    mRightMaterial = CreateForwardLitMaterial(
        "MultiLightRightCube",
        "Tests/Rendering/MultiLightRightCube.material",
        Vector4(0.93f, 0.95f, 0.96f, 1.0f));

    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 MultiLightSceneTest::TearDown() {
    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();

    delete mGroundMaterial;
    mGroundMaterial = nullptr;
    delete mLeftMaterial;
    mLeftMaterial = nullptr;
    delete mCenterMaterial;
    mCenterMaterial = nullptr;
    delete mRightMaterial;
    mRightMaterial = nullptr;
    delete mCubeMesh;
    mCubeMesh = nullptr;

    RHIIntegrationFixture::TearDown();
}

void MultiLightSceneTest::BuildScene() {
    ASSERT_NE(mScene, nullptr);
    ASSERT_NE(mCubeMesh, nullptr);

    GameObject* cameraObject = mScene->CreateGameObject("MainCamera");
    auto* camera = cameraObject->AddComponent<CameraComponent>();
    camera->SetPrimary(true);
    camera->SetFieldOfView(40.0f);
    camera->SetNearClipPlane(0.1f);
    camera->SetFarClipPlane(50.0f);
    camera->SetClearColor(XCEngine::Math::Color(0.025f, 0.025f, 0.035f, 1.0f));
    cameraObject->GetTransform()->SetLocalPosition(Vector3(0.0f, 3.4f, -7.6f));
    cameraObject->GetTransform()->SetLocalRotation(
        Quaternion::LookRotation(Vector3(0.0f, -0.16f, 1.0f).Normalized()));

    GameObject* mainDirectionalObject = mScene->CreateGameObject("MainDirectionalLight");
    auto* mainDirectional = mainDirectionalObject->AddComponent<LightComponent>();
    mainDirectional->SetLightType(LightType::Directional);
    mainDirectional->SetColor(XCEngine::Math::Color(1.0f, 0.97f, 0.93f, 1.0f));
    mainDirectional->SetIntensity(1.25f);
    mainDirectionalObject->GetTransform()->SetLocalRotation(
        Quaternion::LookRotation(Vector3(0.35f, -0.88f, -0.31f).Normalized()));

    GameObject* fillDirectionalObject = mScene->CreateGameObject("FillDirectionalLight");
    auto* fillDirectional = fillDirectionalObject->AddComponent<LightComponent>();
    fillDirectional->SetLightType(LightType::Directional);
    fillDirectional->SetColor(XCEngine::Math::Color(0.22f, 0.55f, 1.0f, 1.0f));
    fillDirectional->SetIntensity(0.42f);
    fillDirectionalObject->GetTransform()->SetLocalRotation(
        Quaternion::LookRotation(Vector3(-0.72f, -0.34f, 0.61f).Normalized()));

    GameObject* pointLightObject = mScene->CreateGameObject("PointLight");
    auto* pointLight = pointLightObject->AddComponent<LightComponent>();
    pointLight->SetLightType(LightType::Point);
    pointLight->SetColor(XCEngine::Math::Color(1.0f, 0.30f, 0.18f, 1.0f));
    pointLight->SetIntensity(6.0f);
    pointLight->SetRange(6.8f);
    pointLightObject->GetTransform()->SetLocalPosition(Vector3(-3.0f, 2.8f, 8.0f));

    GameObject* spotLightObject = mScene->CreateGameObject("SpotLight");
    auto* spotLight = spotLightObject->AddComponent<LightComponent>();
    spotLight->SetLightType(LightType::Spot);
    spotLight->SetColor(XCEngine::Math::Color(0.18f, 1.0f, 0.62f, 1.0f));
    spotLight->SetIntensity(7.0f);
    spotLight->SetRange(11.0f);
    spotLight->SetSpotAngle(38.0f);
    spotLightObject->GetTransform()->SetLocalPosition(Vector3(3.3f, 4.2f, 11.8f));
    spotLightObject->GetTransform()->SetLocalRotation(
        Quaternion::LookRotation(Vector3(-0.18f, -0.84f, -0.52f).Normalized()));

    GameObject* groundObject = mScene->CreateGameObject("Ground");
    auto* groundMeshFilter = groundObject->AddComponent<MeshFilterComponent>();
    auto* groundMeshRenderer = groundObject->AddComponent<MeshRendererComponent>();
    groundObject->GetTransform()->SetLocalPosition(Vector3(0.0f, -0.16f, 10.4f));
    groundObject->GetTransform()->SetLocalScale(Vector3(12.0f, 0.18f, 10.0f));
    groundMeshFilter->SetMesh(ResourceHandle<Mesh>(mCubeMesh));
    groundMeshRenderer->SetMaterial(0, mGroundMaterial);

    GameObject* leftCubeObject = mScene->CreateGameObject("LeftCube");
    auto* leftMeshFilter = leftCubeObject->AddComponent<MeshFilterComponent>();
    auto* leftMeshRenderer = leftCubeObject->AddComponent<MeshRendererComponent>();
    leftCubeObject->GetTransform()->SetLocalPosition(Vector3(-2.8f, 0.80f, 8.3f));
    leftCubeObject->GetTransform()->SetLocalScale(Vector3(1.45f, 1.60f, 1.45f));
    leftMeshFilter->SetMesh(ResourceHandle<Mesh>(mCubeMesh));
    leftMeshRenderer->SetMaterial(0, mLeftMaterial);

    GameObject* centerCubeObject = mScene->CreateGameObject("CenterCube");
    auto* centerMeshFilter = centerCubeObject->AddComponent<MeshFilterComponent>();
    auto* centerMeshRenderer = centerCubeObject->AddComponent<MeshRendererComponent>();
    centerCubeObject->GetTransform()->SetLocalPosition(Vector3(0.0f, 1.0f, 10.4f));
    centerCubeObject->GetTransform()->SetLocalScale(Vector3(1.90f, 2.0f, 1.90f));
    centerMeshFilter->SetMesh(ResourceHandle<Mesh>(mCubeMesh));
    centerMeshRenderer->SetMaterial(0, mCenterMaterial);

    GameObject* rightCubeObject = mScene->CreateGameObject("RightCube");
    auto* rightMeshFilter = rightCubeObject->AddComponent<MeshFilterComponent>();
    auto* rightMeshRenderer = rightCubeObject->AddComponent<MeshRendererComponent>();
    rightCubeObject->GetTransform()->SetLocalPosition(Vector3(2.9f, 0.66f, 12.0f));
    rightCubeObject->GetTransform()->SetLocalScale(Vector3(1.25f, 1.32f, 1.25f));
    rightMeshFilter->SetMesh(ResourceHandle<Mesh>(mCubeMesh));
    rightMeshRenderer->SetMaterial(0, mRightMaterial);
}

RHIResourceView* MultiLightSceneTest::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];
}

void MultiLightSceneTest::RenderFrame() {
    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);
}

TEST_P(MultiLightSceneTest, RenderMultiLightScene) {
    RHICommandQueue* commandQueue = GetCommandQueue();
    RHISwapChain* swapChain = GetSwapChain();
    const int targetFrameCount = 30;
    const char* screenshotFilename = GetScreenshotFilename(GetBackendType());
    const int comparisonThreshold = GetComparisonThreshold(GetBackendType());

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

        BeginRender();
        RenderFrame();

        if (frameCount >= targetFrameCount) {
            commandQueue->WaitForIdle();
            ASSERT_TRUE(TakeScreenshot(screenshotFilename));
            const auto screenshotImage = RenderingIntegrationTestUtils::LoadPpmImage(screenshotFilename);
            ExpectMultiLightKeyPixels(screenshotImage);
            ASSERT_TRUE(CompareWithGoldenTemplate(screenshotFilename, "GT.ppm", static_cast<float>(comparisonThreshold)));
            break;
        }

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

} // namespace

INSTANTIATE_TEST_SUITE_P(D3D12, MultiLightSceneTest, ::testing::Values(RHIType::D3D12));
INSTANTIATE_TEST_SUITE_P(OpenGL, MultiLightSceneTest, ::testing::Values(RHIType::OpenGL));
#if defined(XCENGINE_SUPPORT_VULKAN)
INSTANTIATE_TEST_SUITE_P(Vulkan, MultiLightSceneTest, ::testing::Values(RHIType::Vulkan));
#endif

GTEST_API_ int main(int argc, char** argv) {
    return RunRenderingIntegrationTestMain(argc, argv);
}