#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/Bounds.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/Core/Math/Vector4.h>
#include <XCEngine/Rendering/Planning/CameraRenderRequest.h>
#include <XCEngine/Rendering/RenderContext.h>
#include <XCEngine/Rendering/RenderSurface.h>
#include <XCEngine/Rendering/Execution/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/Resources/Texture/Texture.h>
#include <XCEngine/RHI/RHITexture.h>
#include <XCEngine/Scene/Scene.h>

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

#include <array>
#include <memory>
#include <vector>

using namespace RenderingIntegrationTestUtils;
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 = "alpha_cutout_scene_d3d12.ppm";
constexpr const char* kOpenGLScreenshot = "alpha_cutout_scene_opengl.ppm";
constexpr const char* kVulkanScreenshot = "alpha_cutout_scene_vulkan.ppm";
constexpr uint32_t kFrameWidth = 1280;
constexpr uint32_t kFrameHeight = 720;
constexpr float kAlphaCutoff = 0.5f;

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* CreateVerticalQuadMesh() {
    auto* mesh = new Mesh();
    IResource::ConstructParams params = {};
    params.name = "AlphaCutoutVerticalQuad";
    params.path = "Tests/Rendering/AlphaCutoutVerticalQuad.mesh";
    params.guid = ResourceGUID::Generate(params.path);
    mesh->Initialize(params);

    StaticMeshVertex vertices[4] = {};
    vertices[0].position = Vector3(-0.5f, -0.5f, 0.0f);
    vertices[0].normal = Vector3::Back();
    vertices[0].uv0 = Vector2(0.0f, 1.0f);
    vertices[1].position = Vector3(-0.5f, 0.5f, 0.0f);
    vertices[1].normal = Vector3::Back();
    vertices[1].uv0 = Vector2(0.0f, 0.0f);
    vertices[2].position = Vector3(0.5f, -0.5f, 0.0f);
    vertices[2].normal = Vector3::Back();
    vertices[2].uv0 = Vector2(1.0f, 1.0f);
    vertices[3].position = Vector3(0.5f, 0.5f, 0.0f);
    vertices[3].normal = Vector3::Back();
    vertices[3].uv0 = Vector2(1.0f, 0.0f);

    const uint32_t indices[6] = { 0, 1, 2, 2, 1, 3 };
    mesh->SetVertexData(
        vertices,
        sizeof(vertices),
        4,
        sizeof(StaticMeshVertex),
        VertexAttribute::Position | VertexAttribute::Normal | VertexAttribute::UV0);
    mesh->SetIndexData(indices, sizeof(indices), 6, true);

    const Bounds bounds(Vector3(0.0f, 0.0f, 0.0f), Vector3(1.0f, 1.0f, 0.02f));
    mesh->SetBounds(bounds);

    MeshSection section = {};
    section.baseVertex = 0;
    section.vertexCount = 4;
    section.startIndex = 0;
    section.indexCount = 6;
    section.materialID = 0;
    section.bounds = bounds;
    mesh->AddSection(section);
    return mesh;
}

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

    StaticMeshVertex vertices[4] = {};
    vertices[0].position = Vector3(-0.5f, 0.0f, 0.0f);
    vertices[0].normal = Vector3::Up();
    vertices[0].uv0 = Vector2(0.0f, 1.0f);
    vertices[1].position = Vector3(-0.5f, 0.0f, 1.0f);
    vertices[1].normal = Vector3::Up();
    vertices[1].uv0 = Vector2(0.0f, 0.0f);
    vertices[2].position = Vector3(0.5f, 0.0f, 0.0f);
    vertices[2].normal = Vector3::Up();
    vertices[2].uv0 = Vector2(1.0f, 1.0f);
    vertices[3].position = Vector3(0.5f, 0.0f, 1.0f);
    vertices[3].normal = Vector3::Up();
    vertices[3].uv0 = Vector2(1.0f, 0.0f);

    const uint32_t indices[6] = { 0, 1, 2, 2, 1, 3 };
    mesh->SetVertexData(
        vertices,
        sizeof(vertices),
        4,
        sizeof(StaticMeshVertex),
        VertexAttribute::Position | VertexAttribute::Normal | VertexAttribute::UV0);
    mesh->SetIndexData(indices, sizeof(indices), 6, true);

    const Bounds bounds(Vector3(0.0f, 0.0f, 0.5f), Vector3(1.0f, 0.02f, 1.0f));
    mesh->SetBounds(bounds);

    MeshSection section = {};
    section.baseVertex = 0;
    section.vertexCount = 4;
    section.startIndex = 0;
    section.indexCount = 6;
    section.materialID = 0;
    section.bounds = bounds;
    mesh->AddSection(section);
    return mesh;
}

Texture* CreateCutoutTexture() {
    auto* texture = new Texture();
    IResource::ConstructParams params = {};
    params.name = "AlphaCutoutTexture";
    params.path = "Tests/Rendering/AlphaCutoutScene/cutout.texture";
    params.guid = ResourceGUID::Generate(params.path);
    texture->Initialize(params);

    constexpr uint32_t kTextureSize = 16;
    std::vector<unsigned char> pixels(static_cast<size_t>(kTextureSize) * kTextureSize * 4u, 0u);
    for (uint32_t y = 0; y < kTextureSize; ++y) {
        for (uint32_t x = 0; x < kTextureSize; ++x) {
            const bool insideWindow = x >= 4u && x < 12u && y >= 4u && y < 12u;
            const size_t pixelOffset = (static_cast<size_t>(y) * kTextureSize + x) * 4u;
            pixels[pixelOffset + 0] = 235u;
            pixels[pixelOffset + 1] = 72u;
            pixels[pixelOffset + 2] = 54u;
            pixels[pixelOffset + 3] = insideWindow ? 0u : 255u;
        }
    }

    texture->Create(
        kTextureSize,
        kTextureSize,
        1,
        1,
        XCEngine::Resources::TextureType::Texture2D,
        TextureFormat::RGBA8_UNORM,
        pixels.data(),
        pixels.size());
    return texture;
}

Material* CreateForwardLitMaterial(
    const char* name,
    const char* path,
    const Vector4& baseColor,
    Texture* texture = nullptr,
    bool alphaTest = false) {
    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);

    if (texture != nullptr) {
        material->SetTexture("_MainTex", ResourceHandle<Texture>(texture));
    }

    MaterialRenderState renderState = material->GetRenderState();
    renderState.depthFunc = MaterialComparisonFunc::LessEqual;
    renderState.cullMode = MaterialCullMode::None;
    material->SetRenderState(renderState);

    if (alphaTest) {
        material->SetFloat("_Cutoff", kAlphaCutoff);
        material->EnableKeyword("XC_ALPHA_TEST");
        material->SetRenderQueue(MaterialRenderQueue::AlphaTest);
    } else {
        material->SetRenderQueue(MaterialRenderQueue::Geometry);
    }

    return material;
}

GameObject* CreateRenderable(
    Scene& scene,
    const char* name,
    Mesh* mesh,
    Material* material,
    const Vector3& position,
    const Vector3& scale,
    const Quaternion& rotation = Quaternion::Identity()) {
    GameObject* object = scene.CreateGameObject(name);
    object->GetTransform()->SetLocalPosition(position);
    object->GetTransform()->SetLocalScale(scale);
    object->GetTransform()->SetLocalRotation(rotation);

    auto* meshFilter = object->AddComponent<MeshFilterComponent>();
    auto* meshRenderer = object->AddComponent<MeshRendererComponent>();
    meshFilter->SetMesh(ResourceHandle<Mesh>(mesh));
    meshRenderer->SetMaterial(0, ResourceHandle<Material>(material));
    return object;
}

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 ? 0 : 12;
}

class AlphaCutoutSceneTest : public RHIIntegrationFixture {
protected:
    void SetUp() override;
    void TearDown() override;
    void RenderFrame() override;
    void AssertScenePixels(const char* screenshotFilename);

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

    std::unique_ptr<Scene> mScene;
    std::unique_ptr<SceneRenderer> mSceneRenderer;
    std::vector<RHIResourceView*> mBackBufferViews;
    std::vector<Mesh*> mMeshes;
    std::vector<Material*> mMaterials;
    std::vector<Texture*> mTextures;
    RHITexture* mDepthTexture = nullptr;
    RHIResourceView* mDepthView = nullptr;
};

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

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

    mMeshes.push_back(CreateVerticalQuadMesh());
    mMeshes.push_back(CreateGroundMesh());

    Texture* cutoutTexture = CreateCutoutTexture();
    mTextures.push_back(cutoutTexture);

    mMaterials.push_back(CreateForwardLitMaterial(
        "AlphaCutoutCard",
        "Tests/Rendering/AlphaCutoutScene/card.material",
        Vector4::One(),
        cutoutTexture,
        true));
    mMaterials.push_back(CreateForwardLitMaterial(
        "AlphaCutoutBackdrop",
        "Tests/Rendering/AlphaCutoutScene/backdrop.material",
        Vector4(0.14f, 0.62f, 0.96f, 1.0f)));
    mMaterials.push_back(CreateForwardLitMaterial(
        "AlphaCutoutGround",
        "Tests/Rendering/AlphaCutoutScene/ground.material",
        Vector4(0.92f, 0.93f, 0.95f, 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 AlphaCutoutSceneTest::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();

    for (Material* material : mMaterials) {
        delete material;
    }
    mMaterials.clear();

    for (Texture* texture : mTextures) {
        delete texture;
    }
    mTextures.clear();

    for (Mesh* mesh : mMeshes) {
        delete mesh;
    }
    mMeshes.clear();

    RHIIntegrationFixture::TearDown();
}

void AlphaCutoutSceneTest::BuildScene() {
    ASSERT_NE(mScene, nullptr);
    ASSERT_EQ(mMeshes.size(), 2u);
    ASSERT_EQ(mMaterials.size(), 3u);

    GameObject* cameraObject = mScene->CreateGameObject("MainCamera");
    auto* camera = cameraObject->AddComponent<CameraComponent>();
    camera->SetPrimary(true);
    camera->SetProjectionType(CameraProjectionType::Perspective);
    camera->SetFieldOfView(36.0f);
    camera->SetNearClipPlane(0.1f);
    camera->SetFarClipPlane(40.0f);
    camera->SetClearColor(XCEngine::Math::Color(0.04f, 0.05f, 0.08f, 1.0f));
    cameraObject->GetTransform()->SetLocalPosition(Vector3(0.0f, 1.15f, -1.95f));
    cameraObject->GetTransform()->SetLocalRotation(
        Quaternion::LookRotation(Vector3(0.0f, -0.07f, 1.0f).Normalized()));

    GameObject* lightObject = mScene->CreateGameObject("MainDirectionalLight");
    auto* light = lightObject->AddComponent<LightComponent>();
    light->SetLightType(LightType::Directional);
    light->SetColor(XCEngine::Math::Color(1.0f, 0.98f, 0.95f, 1.0f));
    light->SetIntensity(2.2f);
    light->SetCastsShadows(true);
    lightObject->GetTransform()->SetLocalRotation(
        Quaternion::LookRotation(Vector3(0.45f, -0.78f, 0.44f).Normalized()));

    GameObject* ground = CreateRenderable(
        *mScene,
        "Ground",
        mMeshes[1],
        mMaterials[2],
        Vector3(0.0f, -1.05f, 4.6f),
        Vector3(8.8f, 1.0f, 10.0f));
    auto* groundRenderer = ground->GetComponent<MeshRendererComponent>();
    ASSERT_NE(groundRenderer, nullptr);
    groundRenderer->SetCastShadows(false);
    groundRenderer->SetReceiveShadows(true);

    GameObject* backdrop = CreateRenderable(
        *mScene,
        "Backdrop",
        mMeshes[0],
        mMaterials[1],
        Vector3(0.0f, 1.25f, 9.2f),
        Vector3(5.2f, 4.1f, 1.0f));
    auto* backdropRenderer = backdrop->GetComponent<MeshRendererComponent>();
    ASSERT_NE(backdropRenderer, nullptr);
    backdropRenderer->SetCastShadows(false);
    backdropRenderer->SetReceiveShadows(false);

    GameObject* card = CreateRenderable(
        *mScene,
        "CutoutCard",
        mMeshes[0],
        mMaterials[0],
        Vector3(0.0f, 1.25f, 6.1f),
        Vector3(2.4f, 3.2f, 1.0f));
    auto* cardRenderer = card->GetComponent<MeshRendererComponent>();
    ASSERT_NE(cardRenderer, nullptr);
    cardRenderer->SetCastShadows(true);
    cardRenderer->SetReceiveShadows(false);
}

RHIResourceView* AlphaCutoutSceneTest::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 AlphaCutoutSceneTest::RenderFrame() {
    ASSERT_NE(mScene, nullptr);
    ASSERT_NE(mSceneRenderer, nullptr);

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

    commandList->Reset();

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

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

    std::vector<CameraRenderRequest> requests =
        mSceneRenderer->BuildRenderRequests(*mScene, nullptr, renderContext, mainSurface);
    ASSERT_EQ(requests.size(), 1u);
    ASSERT_TRUE(requests[0].directionalShadow.IsValid());

    RenderSurface depthOnlySurface(kFrameWidth, kFrameHeight);
    depthOnlySurface.SetDepthAttachment(mDepthView);
    depthOnlySurface.SetRenderArea(requests[0].surface.GetRenderArea());
    requests[0].depthOnly.surface = depthOnlySurface;
    requests[0].depthOnly.clearFlags = RenderClearFlags::Depth;
    requests[0].clearFlags = RenderClearFlags::Color;

    ASSERT_TRUE(mSceneRenderer->Render(requests));

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

void AlphaCutoutSceneTest::AssertScenePixels(const char* screenshotFilename) {
    const PpmImage image = LoadPpmImage(screenshotFilename);
    ASSERT_GT(image.width, 0u);
    ASSERT_GT(image.height, 0u);

    ExpectPixelNear(image, 500, 220, { 255, 88, 65 }, 20, "card red frame");
    ExpectPixelNear(image, 640, 260, { 44, 194, 255 }, 20, "cutout hole reveals blue backdrop");
    ExpectPixelLuminanceAtMost(image, 900, 535, 420, "shadow border stays dark");
    ExpectPixelLuminanceAtLeast(image, 800, 535, 700, "shadow hole stays lit");
}

TEST_P(AlphaCutoutSceneTest, RenderAlphaCutoutScene) {
    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));
            AssertScenePixels(screenshotFilename);
            ASSERT_TRUE(CompareWithGoldenTemplate(
                screenshotFilename,
                "GT.ppm",
                static_cast<float>(comparisonThreshold)));
            break;
        }

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

} // namespace

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

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