#include <windows.h>

#include <cstdio>
#include <cstdarg>
#include <cstring>
#include <memory>
#include <vector>

#include "XCEngine/Core/Containers/String.h"
#include "XCEngine/Debug/ConsoleLogSink.h"
#include "XCEngine/Debug/Logger.h"
#include "XCEngine/RHI/RHIBuffer.h"
#include "XCEngine/RHI/RHIEnums.h"
#include "XCEngine/RHI/RHIPipelineState.h"
#include "XCEngine/RHI/RHIResourceView.h"
#include "XCEngine/RHI/Vulkan/VulkanCommandList.h"
#include "XCEngine/RHI/Vulkan/VulkanCommandQueue.h"
#include "XCEngine/RHI/Vulkan/VulkanDevice.h"
#include "XCEngine/RHI/Vulkan/VulkanScreenshot.h"
#include "XCEngine/RHI/Vulkan/VulkanSwapChain.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

using namespace XCEngine::Containers;
using namespace XCEngine::Debug;
using namespace XCEngine::RHI;

namespace {

constexpr int kWidth = 1280;
constexpr int kHeight = 720;
constexpr int kTargetFrameCount = 30;

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 kTriangleVertexShader[] = R"(#version 450
layout(location = 0) in vec4 aPosition;
layout(location = 1) in vec4 aColor;

layout(location = 0) out vec4 vColor;

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

const char kTriangleFragmentShader[] = R"(#version 450
layout(location = 0) in vec4 vColor;
layout(location = 0) out vec4 fragColor;

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

VulkanDevice g_device;
VulkanCommandQueue g_commandQueue;
VulkanSwapChain g_swapChain;
VulkanCommandList g_commandList;
VulkanScreenshot g_screenshot;
std::vector<RHIResourceView*> g_backBufferViews;
RHIBuffer* g_vertexBuffer = nullptr;
RHIResourceView* g_vertexBufferView = nullptr;
RHIBuffer* g_indexBuffer = nullptr;
RHIResourceView* g_indexBufferView = nullptr;
RHIPipelineState* g_pipelineState = nullptr;
HWND g_window = nullptr;

void Log(const char* format, ...) {
    char buffer[1024] = {};
    va_list args;
    va_start(args, format);
    vsnprintf(buffer, sizeof(buffer), format, args);
    va_end(args);
    Logger::Get().Debug(LogCategory::Rendering, String(buffer));
}

LRESULT CALLBACK WindowProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    switch (msg) {
        case WM_CLOSE:
            PostQuitMessage(0);
            return 0;
        default:
            return DefWindowProc(hwnd, msg, wParam, lParam);
    }
}

GraphicsPipelineDesc CreateTrianglePipelineDesc() {
    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);

    desc.vertexShader.source.assign(kTriangleVertexShader, kTriangleVertexShader + std::strlen(kTriangleVertexShader));
    desc.vertexShader.sourceLanguage = ShaderLanguage::GLSL;
    desc.vertexShader.profile = L"vs";

    desc.fragmentShader.source.assign(kTriangleFragmentShader, kTriangleFragmentShader + std::strlen(kTriangleFragmentShader));
    desc.fragmentShader.sourceLanguage = ShaderLanguage::GLSL;
    desc.fragmentShader.profile = L"fs";

    return desc;
}

void ShutdownViews() {
    for (RHIResourceView* view : g_backBufferViews) {
        if (view != nullptr) {
            view->Shutdown();
            delete view;
        }
    }
    g_backBufferViews.clear();
}

void ShutdownTriangleResources() {
    if (g_pipelineState != nullptr) {
        g_pipelineState->Shutdown();
        delete g_pipelineState;
        g_pipelineState = nullptr;
    }

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

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

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

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

void ShutdownVulkan() {
    ShutdownTriangleResources();
    ShutdownViews();
    g_commandList.Shutdown();
    g_swapChain.Shutdown();
    g_commandQueue.Shutdown();
    g_device.Shutdown();
}

bool InitVulkan() {
    RHIDeviceDesc deviceDesc = {};
    deviceDesc.adapterIndex = 0;
    deviceDesc.enableDebugLayer = false;
    deviceDesc.enableGPUValidation = false;

    if (!g_device.Initialize(deviceDesc)) {
        Log("[ERROR] Failed to initialize Vulkan device");
        return false;
    }

    if (!g_commandQueue.Initialize(&g_device, CommandQueueType::Direct)) {
        Log("[ERROR] Failed to initialize Vulkan command queue");
        return false;
    }

    if (!g_swapChain.Initialize(&g_device, &g_commandQueue, g_window, kWidth, kHeight)) {
        Log("[ERROR] Failed to initialize Vulkan swap chain");
        return false;
    }

    if (!g_commandList.Initialize(&g_device)) {
        Log("[ERROR] Failed to initialize Vulkan command list");
        return false;
    }

    Log("[INFO] Vulkan initialized successfully");
    return true;
}

bool InitTriangleResources() {
    BufferDesc vertexBufferDesc = {};
    vertexBufferDesc.size = sizeof(kTriangleVertices);
    vertexBufferDesc.stride = sizeof(Vertex);
    vertexBufferDesc.bufferType = static_cast<uint32_t>(BufferType::Vertex);

    g_vertexBuffer = g_device.CreateBuffer(vertexBufferDesc);
    if (g_vertexBuffer == nullptr) {
        Log("[ERROR] Failed to create vertex buffer");
        return false;
    }
    g_vertexBuffer->SetData(kTriangleVertices, sizeof(kTriangleVertices));
    g_vertexBuffer->SetStride(sizeof(Vertex));
    g_vertexBuffer->SetBufferType(BufferType::Vertex);

    ResourceViewDesc vertexViewDesc = {};
    vertexViewDesc.dimension = ResourceViewDimension::Buffer;
    vertexViewDesc.structureByteStride = sizeof(Vertex);
    g_vertexBufferView = g_device.CreateVertexBufferView(g_vertexBuffer, vertexViewDesc);
    if (g_vertexBufferView == nullptr) {
        Log("[ERROR] Failed to create vertex buffer view");
        return false;
    }

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

    g_indexBuffer = g_device.CreateBuffer(indexBufferDesc);
    if (g_indexBuffer == nullptr) {
        Log("[ERROR] Failed to create index buffer");
        return false;
    }
    g_indexBuffer->SetData(kTriangleIndices, sizeof(kTriangleIndices));
    g_indexBuffer->SetStride(sizeof(uint32_t));
    g_indexBuffer->SetBufferType(BufferType::Index);

    ResourceViewDesc indexViewDesc = {};
    indexViewDesc.dimension = ResourceViewDimension::Buffer;
    indexViewDesc.format = static_cast<uint32_t>(Format::R32_UInt);
    g_indexBufferView = g_device.CreateIndexBufferView(g_indexBuffer, indexViewDesc);
    if (g_indexBufferView == nullptr) {
        Log("[ERROR] Failed to create index buffer view");
        return false;
    }

    GraphicsPipelineDesc pipelineDesc = CreateTrianglePipelineDesc();
    g_pipelineState = g_device.CreatePipelineState(pipelineDesc);
    if (g_pipelineState == nullptr || !g_pipelineState->IsValid()) {
        Log("[ERROR] Failed to create triangle pipeline state");
        return false;
    }

    return true;
}

RHIResourceView* GetCurrentBackBufferView() {
    const uint32_t backBufferIndex = g_swapChain.GetCurrentBackBufferIndex();
    if (g_backBufferViews.size() <= backBufferIndex) {
        g_backBufferViews.resize(backBufferIndex + 1, nullptr);
    }

    if (g_backBufferViews[backBufferIndex] == nullptr) {
        ResourceViewDesc viewDesc = {};
        viewDesc.dimension = ResourceViewDimension::Texture2D;
        viewDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
        viewDesc.arraySize = 1;
        g_backBufferViews[backBufferIndex] = g_device.CreateRenderTargetView(g_swapChain.GetCurrentBackBuffer(), viewDesc);
        if (g_backBufferViews[backBufferIndex] == nullptr) {
            Log("[ERROR] Failed to create render target view for swap chain image %u", backBufferIndex);
        }
    }

    return g_backBufferViews[backBufferIndex];
}

bool RenderFrame() {
    if (!g_swapChain.AcquireNextImage()) {
        Log("[ERROR] Failed to acquire next swap chain image");
        return false;
    }

    RHIResourceView* renderTargetView = GetCurrentBackBufferView();
    if (renderTargetView == nullptr) {
        return false;
    }

    g_commandList.Reset();
    g_commandList.SetRenderTargets(1, &renderTargetView, nullptr);

    Viewport viewport = { 0.0f, 0.0f, static_cast<float>(kWidth), static_cast<float>(kHeight), 0.0f, 1.0f };
    Rect scissorRect = { 0, 0, kWidth, kHeight };
    g_commandList.SetViewport(viewport);
    g_commandList.SetScissorRect(scissorRect);
    g_commandList.Clear(0.0f, 0.0f, 1.0f, 1.0f, 1);

    g_commandList.SetPipelineState(g_pipelineState);
    g_commandList.SetPrimitiveTopology(PrimitiveTopology::TriangleList);

    RHIResourceView* vertexBuffers[] = { g_vertexBufferView };
    uint64_t offsets[] = { 0 };
    uint32_t strides[] = { sizeof(Vertex) };
    g_commandList.SetVertexBuffers(0, 1, vertexBuffers, offsets, strides);
    g_commandList.SetIndexBuffer(g_indexBufferView, 0);
    g_commandList.DrawIndexed(static_cast<uint32_t>(sizeof(kTriangleIndices) / sizeof(kTriangleIndices[0])));
    g_commandList.Close();

    void* commandLists[] = { &g_commandList };
    g_commandQueue.ExecuteCommandLists(1, commandLists);
    return true;
}

} // namespace

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE, LPSTR, int nShowCmd) {
    Logger::Get().Initialize();
    Logger::Get().AddSink(std::make_unique<ConsoleLogSink>());
    Logger::Get().SetMinimumLevel(LogLevel::Debug);

    WNDCLASSEXW wc = {};
    wc.cbSize = sizeof(WNDCLASSEXW);
    wc.style = CS_HREDRAW | CS_VREDRAW;
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = L"XCEngine_Vulkan_Triangle_Test";

    if (!RegisterClassExW(&wc)) {
        Log("[ERROR] Failed to register window class");
        Logger::Get().Shutdown();
        return -1;
    }

    RECT rect = { 0, 0, kWidth, kHeight };
    AdjustWindowRect(&rect, WS_OVERLAPPEDWINDOW, FALSE);

    g_window = CreateWindowExW(
        0,
        L"XCEngine_Vulkan_Triangle_Test",
        L"Vulkan Triangle Integration Test",
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT,
        CW_USEDEFAULT,
        rect.right - rect.left,
        rect.bottom - rect.top,
        nullptr,
        nullptr,
        hInstance,
        nullptr);

    if (g_window == nullptr) {
        Log("[ERROR] Failed to create window");
        Logger::Get().Shutdown();
        return -1;
    }

    if (!InitVulkan() || !InitTriangleResources()) {
        ShutdownVulkan();
        DestroyWindow(g_window);
        g_window = nullptr;
        Logger::Get().Shutdown();
        return -1;
    }

    ShowWindow(g_window, nShowCmd);
    UpdateWindow(g_window);

    MSG msg = {};
    int frameCount = 0;
    int exitCode = 0;

    while (true) {
        if (PeekMessageW(&msg, nullptr, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT) {
                break;
            }
            TranslateMessage(&msg);
            DispatchMessageW(&msg);
            continue;
        }

        if (!RenderFrame()) {
            exitCode = -1;
            break;
        }

        ++frameCount;
        Log("[INFO] Rendered frame %d", frameCount);

        if (frameCount >= kTargetFrameCount) {
            g_commandQueue.WaitForIdle();
            if (!g_screenshot.Capture(&g_device, &g_swapChain, "triangle.ppm")) {
                Log("[ERROR] Failed to capture screenshot");
                exitCode = -1;
            }
            break;
        }

        g_swapChain.Present(0, 0);
    }

    ShutdownVulkan();

    if (g_window != nullptr) {
        DestroyWindow(g_window);
        g_window = nullptr;
    }

    Logger::Get().Shutdown();
    return exitCode;
}