#include <gtest/gtest.h>

#include "../RenderingIntegrationMain.h"

#include <XCEngine/Components/CameraComponent.h>
#include <XCEngine/Components/GameObject.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/Vector2.h>
#include <XCEngine/Core/Math/Vector3.h>
#include <XCEngine/Debug/ConsoleLogSink.h>
#include <XCEngine/Debug/Logger.h>
#include <XCEngine/Rendering/Execution/CameraRenderer.h>
#include <XCEngine/Rendering/Planning/CameraRenderRequest.h>
#include <XCEngine/Rendering/Picking/ObjectIdCodec.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/RHI/RHITexture.h>
#include <XCEngine/Scene/Scene.h>

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

#include <array>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <memory>
#include <string>
#include <vector>
#include <windows.h>

using namespace XCEngine::Components;
using namespace XCEngine::Debug;
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 = "object_id_scene_d3d12.ppm";
constexpr const char* kOpenGLScreenshot = "object_id_scene_opengl.ppm";
constexpr const char* kVulkanScreenshot = "object_id_scene_vulkan.ppm";
constexpr uint32_t kFrameWidth = 1280;
constexpr uint32_t kFrameHeight = 720;

struct PpmImage {
    uint32_t width = 0;
    uint32_t height = 0;
    std::vector<uint8_t> rgb;

    std::array<uint8_t, 3> GetPixel(uint32_t x, uint32_t y) const {
        const size_t index = (static_cast<size_t>(y) * width + x) * 3u;
        return { rgb[index + 0], rgb[index + 1], rgb[index + 2] };
    }
};

std::filesystem::path GetExecutableDirectory() {
    char exePath[MAX_PATH] = {};
    const DWORD length = GetModuleFileNameA(nullptr, exePath, MAX_PATH);
    if (length == 0 || length >= MAX_PATH) {
        return std::filesystem::current_path();
    }

    return std::filesystem::path(exePath).parent_path();
}

std::filesystem::path ResolveRuntimePath(const char* path) {
    std::filesystem::path resolved(path);
    if (resolved.is_absolute()) {
        return resolved;
    }

    return GetExecutableDirectory() / resolved;
}

std::string ReadNextPpmToken(std::istream& stream) {
    std::string token;
    while (stream >> token) {
        if (!token.empty() && token[0] == '#') {
            std::string ignored;
            std::getline(stream, ignored);
            continue;
        }

        return token;
    }

    return {};
}

PpmImage LoadPpmImage(const std::filesystem::path& path) {
    std::ifstream file(path, std::ios::binary);
    EXPECT_TRUE(file.is_open()) << path.string();

    PpmImage image;
    if (!file.is_open()) {
        return image;
    }

    const std::string magic = ReadNextPpmToken(file);
    EXPECT_EQ(magic, "P6");
    if (magic != "P6") {
        return image;
    }

    const std::string widthToken = ReadNextPpmToken(file);
    const std::string heightToken = ReadNextPpmToken(file);
    const std::string maxValueToken = ReadNextPpmToken(file);
    EXPECT_FALSE(widthToken.empty());
    EXPECT_FALSE(heightToken.empty());
    EXPECT_FALSE(maxValueToken.empty());
    if (widthToken.empty() || heightToken.empty() || maxValueToken.empty()) {
        return image;
    }

    image.width = static_cast<uint32_t>(std::stoul(widthToken));
    image.height = static_cast<uint32_t>(std::stoul(heightToken));
    EXPECT_EQ(std::stoul(maxValueToken), 255u);

    file.get();
    image.rgb.resize(static_cast<size_t>(image.width) * image.height * 3u);
    file.read(reinterpret_cast<char*>(image.rgb.data()), static_cast<std::streamsize>(image.rgb.size()));
    EXPECT_EQ(file.gcount(), static_cast<std::streamsize>(image.rgb.size()));
    return image;
}

std::array<uint8_t, 3> EncodeObjectIdToRgb(uint64_t objectId) {
    const uint32_t encodedId = EncodeObjectIdToUInt32(objectId);
    return {
        static_cast<uint8_t>((encodedId >> 0u) & 0xFFu),
        static_cast<uint8_t>((encodedId >> 8u) & 0xFFu),
        static_cast<uint8_t>((encodedId >> 16u) & 0xFFu)
    };
}

void ExpectPixelEquals(
    const PpmImage& image,
    uint32_t x,
    uint32_t y,
    const std::array<uint8_t, 3>& expected,
    const char* label) {
    ASSERT_LT(x, image.width);
    ASSERT_LT(y, image.height);

    const std::array<uint8_t, 3> actual = image.GetPixel(x, y);
    EXPECT_EQ(actual[0], expected[0]) << label << " @ (" << x << ", " << y << ") red";
    EXPECT_EQ(actual[1], expected[1]) << label << " @ (" << x << ", " << y << ") green";
    EXPECT_EQ(actual[2], expected[2]) << label << " @ (" << x << ", " << y << ") blue";
}

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

    StaticMeshVertex vertices[4] = {};
    vertices[0].position = Vector3(-1.0f, -1.0f, 0.0f);
    vertices[0].uv0 = Vector2(0.0f, 1.0f);
    vertices[1].position = Vector3(-1.0f, 1.0f, 0.0f);
    vertices[1].uv0 = Vector2(0.0f, 0.0f);
    vertices[2].position = Vector3(1.0f, -1.0f, 0.0f);
    vertices[2].uv0 = Vector2(1.0f, 1.0f);
    vertices[3].position = Vector3(1.0f, 1.0f, 0.0f);
    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::UV0);
    mesh->SetIndexData(indices, sizeof(indices), 6, true);

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

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

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

Material* CreateUnlitMaterial(Texture* texture) {
    auto* material = new Material();
    IResource::ConstructParams params = {};
    params.name = "ObjectIdMaterial";
    params.path = "Tests/Rendering/ObjectId.material";
    params.guid = ResourceGUID::Generate(params.path);
    material->Initialize(params);
    material->SetShader(ResourceManager::Get().Load<Shader>(GetBuiltinUnlitShaderPath()));
    material->SetShaderPass("Unlit");
    material->SetTexture("_MainTex", ResourceHandle<Texture>(texture));
    return material;
}

GameObject* CreateQuadObject(
    Scene& scene,
    const char* name,
    Mesh* mesh,
    Material* material,
    const Vector3& position,
    const Vector3& scale) {
    GameObject* gameObject = scene.CreateGameObject(name);
    gameObject->GetTransform()->SetLocalPosition(position);
    gameObject->GetTransform()->SetLocalScale(scale);

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

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

class ObjectIdSceneTest : public RHIIntegrationFixture {
protected:
    void SetUp() override;
    void TearDown() override;
    void RenderFrame() override;
    bool RequiresSwapChainAcquire() const override { return false; }
    void VerifyObjectIdScreenshot(const char* screenshotFilename);

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

    std::unique_ptr<Scene> mScene;
    std::unique_ptr<CameraRenderer> mCameraRenderer;
    std::vector<RHIResourceView*> mBackBufferViews;
    RHITexture* mMainColorTexture = nullptr;
    RHIResourceView* mMainColorView = nullptr;
    RHITexture* mObjectIdColorTexture = nullptr;
    RHIResourceView* mObjectIdColorView = nullptr;
    RHITexture* mDepthTexture = nullptr;
    RHIResourceView* mDepthView = nullptr;
    Mesh* mMesh = nullptr;
    Material* mMaterial = nullptr;
    Texture* mTexture = nullptr;
    CameraComponent* mCamera = nullptr;
    GameObject* mLeftObject = nullptr;
    GameObject* mRightObject = nullptr;
    GameObject* mBackObject = nullptr;
    GameObject* mFrontObject = nullptr;
    ResourceStates mBackBufferState = ResourceStates::Present;
};

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

    mCameraRenderer = std::make_unique<CameraRenderer>();
    mScene = std::make_unique<Scene>("ObjectIdScene");
    mMesh = CreateQuadMesh();
    mTexture = CreateSolidTexture();
    mMaterial = CreateUnlitMaterial(mTexture);

    BuildScene();

    TextureDesc colorDesc = {};
    colorDesc.width = kFrameWidth;
    colorDesc.height = kFrameHeight;
    colorDesc.depth = 1;
    colorDesc.mipLevels = 1;
    colorDesc.arraySize = 1;
    colorDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    colorDesc.textureType = static_cast<uint32_t>(XCEngine::RHI::TextureType::Texture2D);
    colorDesc.sampleCount = 1;
    colorDesc.sampleQuality = 0;
    colorDesc.flags = 0;

    mMainColorTexture = GetDevice()->CreateTexture(colorDesc);
    ASSERT_NE(mMainColorTexture, nullptr);
    mObjectIdColorTexture = GetDevice()->CreateTexture(colorDesc);
    ASSERT_NE(mObjectIdColorTexture, nullptr);

    ResourceViewDesc colorViewDesc = {};
    colorViewDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    colorViewDesc.dimension = ResourceViewDimension::Texture2D;
    colorViewDesc.mipLevel = 0;
    mMainColorView = GetDevice()->CreateRenderTargetView(mMainColorTexture, colorViewDesc);
    ASSERT_NE(mMainColorView, nullptr);
    mObjectIdColorView = GetDevice()->CreateRenderTargetView(mObjectIdColorTexture, colorViewDesc);
    ASSERT_NE(mObjectIdColorView, nullptr);

    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);
    mBackBufferState = ResourceStates::Present;
}

void ObjectIdSceneTest::TearDown() {
    mCameraRenderer.reset();

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

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

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

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

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

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

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

    mScene.reset();

    delete mMaterial;
    mMaterial = nullptr;
    delete mMesh;
    mMesh = nullptr;
    delete mTexture;
    mTexture = nullptr;

    mCamera = nullptr;
    mLeftObject = nullptr;
    mRightObject = nullptr;
    mBackObject = nullptr;
    mFrontObject = nullptr;

    RHIIntegrationFixture::TearDown();
}

void ObjectIdSceneTest::BuildScene() {
    ASSERT_NE(mScene, nullptr);
    ASSERT_NE(mMesh, nullptr);
    ASSERT_NE(mMaterial, nullptr);

    GameObject* cameraObject = mScene->CreateGameObject("MainCamera");
    mCamera = cameraObject->AddComponent<CameraComponent>();
    mCamera->SetPrimary(true);
    mCamera->SetProjectionType(CameraProjectionType::Orthographic);
    mCamera->SetOrthographicSize(2.0f);
    mCamera->SetNearClipPlane(0.1f);
    mCamera->SetFarClipPlane(10.0f);
    mCamera->SetClearColor(XCEngine::Math::Color(0.02f, 0.02f, 0.02f, 1.0f));

    mLeftObject = CreateQuadObject(
        *mScene,
        "LeftObject",
        mMesh,
        mMaterial,
        Vector3(-1.8f, 0.0f, 3.0f),
        Vector3(1.1f, 1.1f, 1.0f));

    mRightObject = CreateQuadObject(
        *mScene,
        "RightObject",
        mMesh,
        mMaterial,
        Vector3(1.8f, 0.0f, 3.0f),
        Vector3(1.1f, 1.1f, 1.0f));

    mBackObject = CreateQuadObject(
        *mScene,
        "BackObject",
        mMesh,
        mMaterial,
        Vector3(0.0f, 0.0f, 3.5f),
        Vector3(1.6f, 1.8f, 1.0f));

    mFrontObject = CreateQuadObject(
        *mScene,
        "FrontObject",
        mMesh,
        mMaterial,
        Vector3(0.0f, 0.0f, 2.7f),
        Vector3(0.7f, 0.7f, 1.0f));
}

RHIResourceView* ObjectIdSceneTest::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 ObjectIdSceneTest::RenderFrame() {
    ASSERT_NE(mScene, nullptr);
    ASSERT_NE(mCameraRenderer, nullptr);
    ASSERT_NE(mCamera, nullptr);
    ASSERT_NE(mMainColorView, nullptr);
    ASSERT_NE(mObjectIdColorView, nullptr);
    ASSERT_NE(mDepthView, nullptr);

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

    RHIResourceView* backBufferView = GetCurrentBackBufferView();
    ASSERT_NE(backBufferView, nullptr);

    commandList->Reset();

    RenderSurface mainSurface(kFrameWidth, kFrameHeight);
    mainSurface.SetColorAttachment(mMainColorView);
    mainSurface.SetDepthAttachment(mDepthView);
    mainSurface.SetColorStateBefore(ResourceStates::Common);
    mainSurface.SetColorStateAfter(ResourceStates::Common);

    RenderSurface objectIdSurface(kFrameWidth, kFrameHeight);
    objectIdSurface.SetColorAttachment(mObjectIdColorView);
    objectIdSurface.SetDepthAttachment(mDepthView);
    objectIdSurface.SetColorStateBefore(ResourceStates::Common);
    objectIdSurface.SetColorStateAfter(ResourceStates::CopySrc);

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

    CameraRenderRequest request = {};
    request.scene = mScene.get();
    request.camera = mCamera;
    request.context = renderContext;
    request.surface = mainSurface;
    request.objectId.surface = objectIdSurface;

    ASSERT_TRUE(mCameraRenderer->Render(request));

    commandList->TransitionBarrier(backBufferView, mBackBufferState, ResourceStates::CopyDst);
    commandList->CopyResource(backBufferView, mObjectIdColorView);
    commandList->TransitionBarrier(mObjectIdColorView, ResourceStates::CopySrc, ResourceStates::Common);
    commandList->TransitionBarrier(backBufferView, ResourceStates::CopyDst, ResourceStates::RenderTarget);
    commandList->SetRenderTargets(1, &backBufferView, nullptr);
    mBackBufferState = ResourceStates::RenderTarget;

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

void ObjectIdSceneTest::VerifyObjectIdScreenshot(const char* screenshotFilename) {
    ASSERT_NE(mLeftObject, nullptr);
    ASSERT_NE(mRightObject, nullptr);
    ASSERT_NE(mBackObject, nullptr);
    ASSERT_NE(mFrontObject, nullptr);

    const PpmImage image = LoadPpmImage(ResolveRuntimePath(screenshotFilename));
    ASSERT_EQ(image.width, kFrameWidth);
    ASSERT_EQ(image.height, kFrameHeight);

    ExpectPixelEquals(image, 60, 60, { 0, 0, 0 }, "background");
    ExpectPixelEquals(image, 320, 360, EncodeObjectIdToRgb(mLeftObject->GetID()), "left object");
    ExpectPixelEquals(image, 960, 360, EncodeObjectIdToRgb(mRightObject->GetID()), "right object");
    ExpectPixelEquals(image, 640, 220, EncodeObjectIdToRgb(mBackObject->GetID()), "back object visible region");
    ExpectPixelEquals(image, 640, 360, EncodeObjectIdToRgb(mFrontObject->GetID()), "front object overlap region");
}

TEST_P(ObjectIdSceneTest, RenderObjectIdScene) {
    RHICommandQueue* commandQueue = GetCommandQueue();
    const int targetFrameCount = 30;
    const char* screenshotFilename = GetScreenshotFilename(GetBackendType());

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

        BeginRender();
        RenderFrame();

        if (frameCount >= targetFrameCount) {
            commandQueue->WaitForIdle();
            ASSERT_TRUE(TakeScreenshot(screenshotFilename));
            VerifyObjectIdScreenshot(screenshotFilename);
            break;
        }
    }
}

} // namespace

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

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