#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <filesystem>
#include <fstream>

#include <gtest/gtest.h>

#include "../fixtures/RHIIntegrationFixture.h"
#include "XCEngine/Debug/ConsoleLogSink.h"
#include "XCEngine/Debug/Logger.h"
#include "XCEngine/RHI/RHIBuffer.h"
#include "XCEngine/RHI/RHIPipelineState.h"
#include "XCEngine/RHI/RHIResourceView.h"

using namespace XCEngine::Debug;
using namespace XCEngine::RHI;
using namespace XCEngine::RHI::Integration;

namespace {

struct Vertex {
    float pos[4];
    float col[4];
};

constexpr Vertex kTriangleVertices[] = {
    { {  0.0f,  0.5f, 0.0f, 1.0f }, { 1.0f, 0.0f, 0.0f, 1.0f } },
    { { -0.5f, -0.5f, 0.0f, 1.0f }, { 0.0f, 1.0f, 0.0f, 1.0f } },
    { {  0.5f, -0.5f, 0.0f, 1.0f }, { 0.0f, 0.0f, 1.0f, 1.0f } },
};

constexpr uint32_t kTriangleIndices[] = { 0, 1, 2 };

const char kTriangleHlsl[] = R"(
struct VSInput {
    float4 position : POSITION;
    float4 color : COLOR;
};

struct PSInput {
    float4 position : SV_POSITION;
    float4 color : COLOR;
};

PSInput MainVS(VSInput input) {
    PSInput output;
    output.position = input.position;
    output.color = input.color;
    return output;
}

float4 MainPS(PSInput input) : SV_TARGET {
    return input.color;
}
)";

const char kTriangleVertexShader[] = R"(#version 430
layout(location = 0) in vec4 aPosition;
layout(location = 1) in vec4 aColor;

out vec4 vColor;

void main() {
    gl_Position = aPosition;
    vColor = aColor;
}
)";

const char kTriangleFragmentShader[] = R"(#version 430
in vec4 vColor;

layout(location = 0) out vec4 fragColor;

void main() {
    fragColor = vColor;
}
)";

std::filesystem::path GetExecutableDirectory() {
    char exePath[MAX_PATH] = {};
    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::vector<uint8_t> LoadBinaryFileRelative(const char* filename) {
    const std::filesystem::path path = GetExecutableDirectory() / filename;
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    if (!file.is_open()) {
        return {};
    }

    const std::streamsize size = file.tellg();
    if (size <= 0) {
        return {};
    }

    std::vector<uint8_t> bytes(static_cast<size_t>(size));
    file.seekg(0, std::ios::beg);
    if (!file.read(reinterpret_cast<char*>(bytes.data()), size)) {
        return {};
    }

    return bytes;
}

const char* GetScreenshotFilename(RHIType type) {
    switch (type) {
        case RHIType::D3D12:
            return "triangle_d3d12.ppm";
        case RHIType::OpenGL:
            return "triangle_opengl.ppm";
        case RHIType::Vulkan:
            return "triangle_vulkan.ppm";
        default:
            return "triangle_unknown.ppm";
    }
}

int GetComparisonThreshold(RHIType type) {
    return type == RHIType::OpenGL ? 5 : 0;
}

GraphicsPipelineDesc CreateTrianglePipelineDesc(RHIType type) {
    GraphicsPipelineDesc desc = {};
    desc.topologyType = static_cast<uint32_t>(PrimitiveTopologyType::Triangle);
    desc.renderTargetFormats[0] = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    desc.depthStencilFormat = static_cast<uint32_t>(Format::Unknown);
    desc.sampleCount = 1;

    desc.rasterizerState.fillMode = static_cast<uint32_t>(FillMode::Solid);
    desc.rasterizerState.cullMode = static_cast<uint32_t>(CullMode::None);
    desc.rasterizerState.frontFace = static_cast<uint32_t>(FrontFace::CounterClockwise);
    desc.rasterizerState.depthClipEnable = true;

    desc.depthStencilState.depthTestEnable = false;
    desc.depthStencilState.depthWriteEnable = false;
    desc.depthStencilState.stencilEnable = false;

    InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(Format::R32G32B32A32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = 0;
    desc.inputLayout.elements.push_back(position);

    InputElementDesc color = {};
    color.semanticName = "COLOR";
    color.semanticIndex = 0;
    color.format = static_cast<uint32_t>(Format::R32G32B32A32_Float);
    color.inputSlot = 0;
    color.alignedByteOffset = sizeof(float) * 4;
    desc.inputLayout.elements.push_back(color);

    if (type == RHIType::D3D12) {
        desc.vertexShader.source.assign(kTriangleHlsl, kTriangleHlsl + strlen(kTriangleHlsl));
        desc.vertexShader.sourceLanguage = ShaderLanguage::HLSL;
        desc.vertexShader.entryPoint = L"MainVS";
        desc.vertexShader.profile = L"vs_5_0";

        desc.fragmentShader.source.assign(kTriangleHlsl, kTriangleHlsl + strlen(kTriangleHlsl));
        desc.fragmentShader.sourceLanguage = ShaderLanguage::HLSL;
        desc.fragmentShader.entryPoint = L"MainPS";
        desc.fragmentShader.profile = L"ps_5_0";
    } else if (type == RHIType::OpenGL) {
        desc.vertexShader.source.assign(kTriangleVertexShader, kTriangleVertexShader + strlen(kTriangleVertexShader));
        desc.vertexShader.sourceLanguage = ShaderLanguage::GLSL;
        desc.vertexShader.profile = L"vs_4_30";

        desc.fragmentShader.source.assign(kTriangleFragmentShader, kTriangleFragmentShader + strlen(kTriangleFragmentShader));
        desc.fragmentShader.sourceLanguage = ShaderLanguage::GLSL;
        desc.fragmentShader.profile = L"fs_4_30";
    } else if (type == RHIType::Vulkan) {
        desc.vertexShader.source = LoadBinaryFileRelative("triangle_vulkan.vert.spv");
        desc.vertexShader.sourceLanguage = ShaderLanguage::SPIRV;
        desc.vertexShader.entryPoint = L"main";

        desc.fragmentShader.source = LoadBinaryFileRelative("triangle_vulkan.frag.spv");
        desc.fragmentShader.sourceLanguage = ShaderLanguage::SPIRV;
        desc.fragmentShader.entryPoint = L"main";
    }

    return desc;
}

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

private:
    void InitializeTriangleResources();
    void ShutdownTriangleResources();

    RHIBuffer* mVertexBuffer = nullptr;
    RHIResourceView* mVertexBufferView = nullptr;
    RHIBuffer* mIndexBuffer = nullptr;
    RHIResourceView* mIndexBufferView = nullptr;
    RHIPipelineState* mPipelineState = nullptr;
};

void TriangleTest::SetUp() {
    RHIIntegrationFixture::SetUp();
    InitializeTriangleResources();
}

void TriangleTest::TearDown() {
    ShutdownTriangleResources();
    RHIIntegrationFixture::TearDown();
}

void TriangleTest::InitializeTriangleResources() {
    BufferDesc bufferDesc = {};
    bufferDesc.size = sizeof(kTriangleVertices);
    bufferDesc.stride = sizeof(Vertex);
    bufferDesc.bufferType = static_cast<uint32_t>(BufferType::Vertex);

    mVertexBuffer = GetDevice()->CreateBuffer(bufferDesc);
    ASSERT_NE(mVertexBuffer, nullptr);
    mVertexBuffer->SetData(kTriangleVertices, sizeof(kTriangleVertices));
    mVertexBuffer->SetStride(sizeof(Vertex));
    mVertexBuffer->SetBufferType(BufferType::Vertex);

    ResourceViewDesc viewDesc = {};
    viewDesc.dimension = ResourceViewDimension::Buffer;
    viewDesc.structureByteStride = sizeof(Vertex);
    mVertexBufferView = GetDevice()->CreateVertexBufferView(mVertexBuffer, viewDesc);
    ASSERT_NE(mVertexBufferView, nullptr);

    BufferDesc indexBufferDesc = {};
    indexBufferDesc.size = sizeof(kTriangleIndices);
    indexBufferDesc.stride = sizeof(uint32_t);
    indexBufferDesc.bufferType = static_cast<uint32_t>(BufferType::Index);

    mIndexBuffer = GetDevice()->CreateBuffer(indexBufferDesc);
    ASSERT_NE(mIndexBuffer, nullptr);
    mIndexBuffer->SetData(kTriangleIndices, sizeof(kTriangleIndices));
    mIndexBuffer->SetStride(sizeof(uint32_t));
    mIndexBuffer->SetBufferType(BufferType::Index);

    ResourceViewDesc indexViewDesc = {};
    indexViewDesc.dimension = ResourceViewDimension::Buffer;
    indexViewDesc.format = static_cast<uint32_t>(Format::R32_UInt);
    mIndexBufferView = GetDevice()->CreateIndexBufferView(mIndexBuffer, indexViewDesc);
    ASSERT_NE(mIndexBufferView, nullptr);

    GraphicsPipelineDesc pipelineDesc = CreateTrianglePipelineDesc(GetBackendType());
    mPipelineState = GetDevice()->CreatePipelineState(pipelineDesc);
    ASSERT_NE(mPipelineState, nullptr);
    ASSERT_TRUE(mPipelineState->IsValid());

    Log("[TEST] Triangle resources initialized for backend %d", static_cast<int>(GetBackendType()));
}

void TriangleTest::ShutdownTriangleResources() {
    if (mPipelineState != nullptr) {
        mPipelineState->Shutdown();
        delete mPipelineState;
        mPipelineState = nullptr;
    }

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

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

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

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

void TriangleTest::RenderFrame() {
    ASSERT_NE(mPipelineState, nullptr);
    ASSERT_NE(mVertexBufferView, nullptr);
    ASSERT_NE(mIndexBufferView, nullptr);

    RHICommandList* cmdList = GetCommandList();
    RHICommandQueue* cmdQueue = GetCommandQueue();
    ASSERT_NE(cmdList, nullptr);
    ASSERT_NE(cmdQueue, nullptr);

    cmdList->Reset();
    SetRenderTargetForClear();

    Viewport viewport = { 0.0f, 0.0f, 1280.0f, 720.0f, 0.0f, 1.0f };
    Rect scissorRect = { 0, 0, 1280, 720 };
    cmdList->SetViewport(viewport);
    cmdList->SetScissorRect(scissorRect);
    cmdList->Clear(0.0f, 0.0f, 1.0f, 1.0f, 1);

    cmdList->SetPipelineState(mPipelineState);
    cmdList->SetPrimitiveTopology(PrimitiveTopology::TriangleList);

    RHIResourceView* vertexBuffers[] = { mVertexBufferView };
    uint64_t offsets[] = { 0 };
    uint32_t strides[] = { sizeof(Vertex) };
    cmdList->SetVertexBuffers(0, 1, vertexBuffers, offsets, strides);
    cmdList->SetIndexBuffer(mIndexBufferView, 0);
    cmdList->DrawIndexed(static_cast<uint32_t>(sizeof(kTriangleIndices) / sizeof(kTriangleIndices[0])));

    EndRender();

    cmdList->Close();
    void* cmdLists[] = { cmdList };
    cmdQueue->ExecuteCommandLists(1, cmdLists);
}

TEST_P(TriangleTest, RenderTriangle) {
    RHICommandQueue* cmdQueue = 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) {
            cmdQueue->WaitForPreviousFrame();
        }

        Log("[TEST] Triangle MainLoop: frame %d", frameCount);
        BeginRender();
        RenderFrame();

        if (frameCount >= targetFrameCount) {
            cmdQueue->WaitForIdle();
            Log("[TEST] Triangle MainLoop: frame %d reached, capturing %s", frameCount, screenshotFilename);
            ASSERT_TRUE(TakeScreenshot(screenshotFilename));
            ASSERT_TRUE(CompareWithGoldenTemplate(screenshotFilename, "GT.ppm", static_cast<float>(comparisonThreshold)));
            Log("[TEST] Triangle MainLoop: frame %d compare passed", frameCount);
            break;
        }

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

} // namespace

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

GTEST_API_ int main(int argc, char** argv) {
    Logger::Get().Initialize();
    Logger::Get().AddSink(std::make_unique<ConsoleLogSink>());
    Logger::Get().SetMinimumLevel(LogLevel::Debug);

    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}