#include <windows.h>
#include <d3d12.h>
#include <dxgi1_4.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <stdarg.h>
#include <string>

#include "XCEngine/RHI/RHIEnums.h"
#include "XCEngine/RHI/RHITypes.h"
#include "XCEngine/RHI/D3D12/D3D12Device.h"
#include "XCEngine/RHI/D3D12/D3D12CommandQueue.h"
#include "XCEngine/RHI/D3D12/D3D12CommandAllocator.h"
#include "XCEngine/RHI/D3D12/D3D12CommandList.h"
#include "XCEngine/RHI/D3D12/D3D12DescriptorHeap.h"
#include "XCEngine/RHI/D3D12/D3D12Fence.h"
#include "XCEngine/RHI/D3D12/D3D12SwapChain.h"
#include "XCEngine/RHI/D3D12/D3D12Buffer.h"
#include "XCEngine/RHI/D3D12/D3D12Texture.h"
#include "XCEngine/RHI/D3D12/D3D12ResourceView.h"
#include "XCEngine/RHI/D3D12/D3D12Screenshot.h"
#include "XCEngine/Debug/Logger.h"
#include "XCEngine/Debug/ConsoleLogSink.h"
#include "XCEngine/Debug/FileLogSink.h"
#include "XCEngine/Debug/RenderDocCapture.h"
#include "XCEngine/Containers/String.h"

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

#pragma comment(lib,"d3d12.lib")
#pragma comment(lib,"dxgi.lib")
#pragma comment(lib,"dxguid.lib")
#pragma comment(lib,"d3dcompiler.lib")
#pragma comment(lib,"winmm.lib")

// Global D3D12 objects
D3D12Device gDevice;
D3D12CommandQueue gCommandQueue;
D3D12SwapChain gSwapChain;
D3D12CommandAllocator gCommandAllocator;
D3D12CommandList gCommandList;

// Render targets
D3D12Texture gDepthStencil;
D3D12DescriptorHeap gRTVHeap;
D3D12DescriptorHeap gDSVHeap;
D3D12ResourceView gRTVs[2];
D3D12ResourceView gDSV;

UINT gRTVDescriptorSize = 0;
UINT gDSVDescriptorSize = 0;
int gCurrentRTIndex = 0;

// Window
HWND gHWND = nullptr;
int gWidth = 1280;
int gHeight = 720;

// Log helper
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));
}

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

// Initialize D3D12
bool InitD3D12() {
    // Create device
    RHIDeviceDesc deviceDesc;
    deviceDesc.windowHandle = gHWND;
    deviceDesc.width = gWidth;
    deviceDesc.height = gHeight;
    deviceDesc.adapterIndex = 0;
    deviceDesc.enableDebugLayer = false;
    deviceDesc.enableGPUValidation = false;

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

    ID3D12Device* device = gDevice.GetDevice();
    IDXGIFactory4* factory = gDevice.GetFactory();

    // Create command queue
    if (!gCommandQueue.Initialize(device, CommandQueueType::Direct)) {
        Log("[ERROR] Failed to initialize command queue");
        return false;
    }

    // Create swap chain using encapsulated interface
    if (!gSwapChain.Initialize(factory, gCommandQueue.GetCommandQueue(), gHWND, gWidth, gHeight, 2)) {
        Log("[ERROR] Failed to initialize swap chain");
        return false;
    }

    // Initialize depth stencil
    gDepthStencil.InitializeDepthStencil(device, gWidth, gHeight);

    // Create RTV heap
    gRTVHeap.Initialize(device, DescriptorHeapType::RTV, 2);
    gRTVDescriptorSize = gDevice.GetDescriptorHandleIncrementSize(DescriptorHeapType::RTV);

    // Create DSV heap
    gDSVHeap.Initialize(device, DescriptorHeapType::DSV, 1);
    gDSVDescriptorSize = gDevice.GetDescriptorHandleIncrementSize(DescriptorHeapType::DSV);

    // Create RTVs for back buffers using encapsulated interface
    for (int i = 0; i < 2; i++) {
        D3D12Texture& backBuffer = gSwapChain.GetBackBuffer(i);

        CPUDescriptorHandle rtvCpuHandle = gRTVHeap.GetCPUDescriptorHandle(i);
        D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle = { rtvCpuHandle.ptr };
        D3D12_RENDER_TARGET_VIEW_DESC rtvDesc = D3D12ResourceView::CreateRenderTargetDesc(Format::R8G8B8A8_UNorm, D3D12_RTV_DIMENSION_TEXTURE2D);
        gRTVs[i].InitializeAsRenderTarget(device, backBuffer.GetResource(), &rtvDesc, &gRTVHeap, i);
    }

    // Create DSV
    D3D12_DEPTH_STENCIL_VIEW_DESC dsvDesc = D3D12ResourceView::CreateDepthStencilDesc(Format::D24_UNorm_S8_UInt, D3D12_DSV_DIMENSION_TEXTURE2D);
    gDSV.InitializeAsDepthStencil(device, gDepthStencil.GetResource(), &dsvDesc, &gDSVHeap, 0);

    // Create command allocator and list
    gCommandAllocator.Initialize(device, CommandQueueType::Direct);
    gCommandList.Initialize(device, CommandQueueType::Direct, gCommandAllocator.GetCommandAllocator());

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

// Wait for GPU
void WaitForGPU() {
    gCommandQueue.WaitForIdle();
}

// Execute command list
void ExecuteCommandList() {
    gCommandList.Close();
    void* commandLists[] = { &gCommandList };
    gCommandQueue.ExecuteCommandLists(1, commandLists);
}

// Begin rendering
void BeginRender() {
    gCurrentRTIndex = gSwapChain.GetCurrentBackBufferIndex();

    // Transition render target
    D3D12Texture& currentBackBuffer = gSwapChain.GetBackBuffer(gCurrentRTIndex);
    gCommandList.TransitionBarrier(currentBackBuffer.GetResource(),
        ResourceStates::Present, ResourceStates::RenderTarget);

    // Set render targets using encapsulated interface
    CPUDescriptorHandle rtvCpuHandle = gRTVHeap.GetCPUDescriptorHandle(gCurrentRTIndex);
    CPUDescriptorHandle dsvCpuHandle = gDSVHeap.GetCPUDescriptorHandle(0);
    D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle = { rtvCpuHandle.ptr };
    D3D12_CPU_DESCRIPTOR_HANDLE dsvHandle = { dsvCpuHandle.ptr };

    gCommandList.SetRenderTargetsHandle(1, &rtvHandle, &dsvHandle);

    // Set viewport and scissor
    Viewport viewport = { 0.0f, 0.0f, (float)gWidth, (float)gHeight, 0.0f, 1.0f };
    Rect scissorRect = { 0, 0, gWidth, gHeight };
    gCommandList.SetViewport(viewport);
    gCommandList.SetScissorRect(scissorRect);

    // Clear
    float clearColor[] = { 1.0f, 0.0f, 0.0f, 1.0f };
    gCommandList.ClearRenderTargetView(rtvHandle, clearColor, 0, nullptr);
    gCommandList.ClearDepthStencilView(dsvHandle, D3D12_CLEAR_FLAG_DEPTH | D3D12_CLEAR_FLAG_STENCIL, 1.0f, 0, 0, nullptr);
}

// End rendering
void EndRender() {
    D3D12Texture& currentBackBuffer = gSwapChain.GetBackBuffer(gCurrentRTIndex);
    gCommandList.TransitionBarrier(currentBackBuffer.GetResource(),
        ResourceStates::RenderTarget, ResourceStates::Present);
}

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

    Log("[INFO] D3D12 Integration Test Starting");

    // Register window class
    WNDCLASSEX wc = {};
    wc.cbSize = sizeof(WNDCLASSEX);
    wc.style = CS_HREDRAW | CS_VREDRAW;
    wc.lpfnWndProc = WindowProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = L"D3D12Test";

    if (!RegisterClassEx(&wc)) {
        MessageBox(NULL, L"Failed to register window class", L"Error", MB_OK);
        return -1;
    }

    // Create window
    RECT rect = { 0, 0, gWidth, gHeight };
    AdjustWindowRect(&rect, WS_OVERLAPPEDWINDOW, FALSE);

    gHWND = CreateWindowEx(0, L"D3D12Test", L"D3D12 Integration Test",
        WS_OVERLAPPEDWINDOW, CW_USEDEFAULT, CW_USEDEFAULT,
        rect.right - rect.left, rect.bottom - rect.top,
        NULL, NULL, hInstance, NULL);

    if (!gHWND) {
        MessageBox(NULL, L"Failed to create window", L"Error", MB_OK);
        return -1;
    }

    RenderDocCapture::Get().Initialize(nullptr, gHWND);
    RenderDocCapture::Get().SetCaptureFilePath(".\\minimal_frame30");

    // Initialize D3D12
    if (!InitD3D12()) {
        MessageBox(NULL, L"Failed to initialize D3D12", L"Error", MB_OK);
        return -1;
    }

    // Set device for RenderDoc (must be called after D3D12 init)
    RenderDocCapture::Get().SetDevice(gDevice.GetDevice());

    // Show window
    ShowWindow(gHWND, nShowCmd);
    UpdateWindow(gHWND);

    // Main loop
    MSG msg = {};
    int frameCount = 0;
    const int targetFrameCount = 30;

    while (true) {
        if (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT) {
                break;
            }
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        } else {
            // Wait for previous frame to complete before resetting
            if (frameCount > 0) {
                gCommandQueue.WaitForPreviousFrame();
            }

            // Reset command list for this frame
            gCommandAllocator.Reset();
            gCommandList.Reset();

            // Render
            BeginRender();
            // (Add rendering code here)
            
            frameCount++;

            EndRender();

            // Execute
            ExecuteCommandList();

            if (frameCount >= targetFrameCount) {
                if (RenderDocCapture::Get().EndCapture()) {
                    Log("[INFO] RenderDoc capture ended");
                }
                WaitForGPU();
                Log("[INFO] GPU idle, taking screenshot...");
                bool screenshotResult = D3D12Screenshot::Capture(
                    gDevice,
                    gCommandQueue,
                    gSwapChain.GetBackBuffer(gCurrentRTIndex),
                    "minimal.ppm"
                );
                if (screenshotResult) {
                    Log("[INFO] Screenshot saved to minimal.ppm");
                } else {
                    Log("[ERROR] Screenshot failed");
                }
                Log("[INFO] RenderDoc capture completed");
                break;
            }

            if (frameCount == targetFrameCount - 1) {
                if (RenderDocCapture::Get().BeginCapture("D3D12_Minimal_Test")) {
                    Log("[INFO] RenderDoc capture started");
                }
            }

            EndRender();
            ExecuteCommandList();
            gSwapChain.Present(0, 0);

            if (frameCount >= targetFrameCount) {
                WaitForGPU();
                Log("[INFO] GPU idle, taking screenshot...");
                bool screenshotResult = D3D12Screenshot::Capture(
                    gDevice,
                    gCommandQueue,
                    gSwapChain.GetBackBuffer(gCurrentRTIndex),
                    "minimal.ppm"
                );
                if (screenshotResult) {
                    Log("[INFO] Screenshot saved to minimal.ppm");
                } else {
                    Log("[ERROR] Screenshot failed");
                }
                Log("[INFO] RenderDoc capture completed");
                break;
            }
        }
    }

    // Shutdown
    RenderDocCapture::Get().Shutdown();
    gCommandList.Shutdown();
    gCommandAllocator.Shutdown();
    gSwapChain.Shutdown();
    gDevice.Shutdown();

    Logger::Get().Shutdown();

    Log("[INFO] D3D12 Integration Test Finished");
    return 0;
}