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

#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/D3D12Shader.h"
#include "XCEngine/RHI/D3D12/D3D12RootSignature.h"
#include "XCEngine/RHI/D3D12/D3D12PipelineState.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"
#include "XCEngine/Math/Matrix4.h"
#include "XCEngine/Math/Vector3.h"
#include "third_party/stb/stb_image.h"

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

#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")

D3D12Device gDevice;
D3D12CommandQueue gCommandQueue;
D3D12SwapChain gSwapChain;
D3D12CommandAllocator gCommandAllocator;
D3D12CommandList gCommandList;
D3D12Texture gDepthStencil;
D3D12DescriptorHeap gRTVHeap;
D3D12DescriptorHeap gDSVHeap;
D3D12DescriptorHeap gSRVHeap;
D3D12ResourceView gRTVs[2];
D3D12ResourceView gDSV;

D3D12Shader gVertexShader;
D3D12Shader gPixelShader;
D3D12RootSignature gRootSignature;
D3D12PipelineState gPipelineState;
D3D12Buffer gVertexBuffer;
D3D12Buffer gIndexBuffer;
D3D12Buffer gMVPBuffer;
D3D12Texture gDiffuseTexture;
D3D12ResourceView gDiffuseSRV;

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

UINT gIndexCount = 0;

Matrix4x4 gProjectionMatrix;
Matrix4x4 gViewMatrix;
Matrix4x4 gModelMatrix;

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

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);
        break;
    }
    return DefWindowProc(hwnd, msg, wParam, lParam);
}

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

void GenerateSphere(std::vector<Vertex>& vertices, std::vector<UINT32>& indices, float radius, int segments) {
    vertices.clear();
    indices.clear();

    segments = (segments < 3) ? 3 : segments;

    for (int lat = 0; lat <= segments; ++lat) {
        float phi = static_cast<float>(M_PI) * lat / segments;
        float sinPhi = sinf(phi);
        float cosPhi = cosf(phi);

        for (int lon = 0; lon <= segments; ++lon) {
            float theta = static_cast<float>(2 * M_PI) * lon / segments;
            float sinTheta = sinf(theta);
            float cosTheta = cosf(theta);

            Vertex v;
            v.pos[0] = radius * sinPhi * cosTheta;
            v.pos[1] = radius * cosPhi;
            v.pos[2] = radius * sinPhi * sinTheta;
            v.pos[3] = 1.0f;

            v.texcoord[0] = static_cast<float>(lon) / segments;
            v.texcoord[1] = static_cast<float>(lat) / segments;
            v.texcoord[2] = 0.0f;
            v.texcoord[3] = 0.0f;

            vertices.push_back(v);
        }
    }

    for (int lat = 0; lat < segments; ++lat) {
        for (int lon = 0; lon < segments; ++lon) {
            UINT32 topLeft = lat * (segments + 1) + lon;
            UINT32 topRight = topLeft + 1;
            UINT32 bottomLeft = (lat + 1) * (segments + 1) + lon;
            UINT32 bottomRight = bottomLeft + 1;

            indices.push_back(topLeft);
            indices.push_back(bottomLeft);
            indices.push_back(topRight);

            indices.push_back(topRight);
            indices.push_back(bottomLeft);
            indices.push_back(bottomRight);
        }
    }
}

bool LoadTexture(const char* filename, D3D12Texture& texture, D3D12ResourceView& srv, ID3D12Device* device, D3D12DescriptorHeap& srvHeap, ID3D12GraphicsCommandList* commandList, D3D12CommandAllocator& allocator, D3D12CommandQueue& queue) {
    int width, height, channels;
    stbi_uc* pixels = stbi_load(filename, &width, &height, &channels, STBI_rgb_alpha);
    if (!pixels) {
        Log("[ERROR] Failed to load texture: %s", filename);
        return false;
    }

    allocator.Reset();
    commandList->Reset(allocator.GetCommandAllocator(), nullptr);

    if (!texture.InitializeFromData(device, commandList, pixels, width, height, DXGI_FORMAT_R8G8B8A8_UNORM)) {
        Log("[ERROR] Failed to initialize texture");
        stbi_image_free(pixels);
        return false;
    }

    commandList->Close();
    ID3D12CommandList* lists[] = { commandList };
    queue.ExecuteCommandListsInternal(1, lists);
    queue.WaitForIdle();

    texture.SetName(filename);

    srvHeap.Initialize(device, DescriptorHeapType::CBV_SRV_UAV, 1, true);
    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = D3D12ResourceView::CreateShaderResourceDesc(Format::R8G8B8A8_UNorm, D3D12_SRV_DIMENSION_TEXTURE2D);
    srv.InitializeAsShaderResource(device, texture.GetResource(), &srvDesc, &srvHeap, 0);

    return true;
}

bool InitD3D12() {
    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();

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

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

    gDepthStencil.InitializeDepthStencil(device, gWidth, gHeight);

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

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

    for (int i = 0; i < 2; i++) {
        D3D12Texture& backBuffer = gSwapChain.GetBackBuffer(i);
        D3D12_RENDER_TARGET_VIEW_DESC rtvDesc = D3D12ResourceView::CreateRenderTargetDesc(Format::R8G8B8A8_UNorm, D3D12_RTV_DIMENSION_TEXTURE2D);
        gRTVs[i].InitializeAsRenderTarget(device, backBuffer.GetResource(), &rtvDesc, &gRTVHeap, i);
    }

    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);

    gCommandAllocator.Initialize(device, CommandQueueType::Direct);
    gCommandList.Initialize(device, CommandQueueType::Direct, gCommandAllocator.GetCommandAllocator());

    if (!gVertexShader.CompileFromFile(L"Res/Shader/sphere.hlsl", "MainVS", "vs_5_1")) {
        Log("[ERROR] Failed to compile vertex shader");
        return false;
    }
    Log("[INFO] Vertex shader compiled, bytecode size: %zu", gVertexShader.GetBytecodeSize());

    if (!gPixelShader.CompileFromFile(L"Res/Shader/sphere.hlsl", "MainPS", "ps_5_1")) {
        Log("[ERROR] Failed to compile pixel shader");
        return false;
    }
    Log("[INFO] Pixel shader compiled, bytecode size: %zu", gPixelShader.GetBytecodeSize());

    D3D12_DESCRIPTOR_RANGE descriptorRange = D3D12RootSignature::CreateDescriptorRange(
        D3D12_DESCRIPTOR_RANGE_TYPE_SRV, 0, 1, 0);

    D3D12_ROOT_PARAMETER rootParameters[2];
    rootParameters[0].ParameterType = D3D12_ROOT_PARAMETER_TYPE_CBV;
    rootParameters[0].ShaderVisibility = D3D12_SHADER_VISIBILITY_VERTEX;
    rootParameters[0].Descriptor.ShaderRegister = 0;
    rootParameters[0].Descriptor.RegisterSpace = 0;

    rootParameters[1].ParameterType = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
    rootParameters[1].ShaderVisibility = D3D12_SHADER_VISIBILITY_PIXEL;
    rootParameters[1].DescriptorTable.NumDescriptorRanges = 1;
    rootParameters[1].DescriptorTable.pDescriptorRanges = &descriptorRange;

    D3D12_STATIC_SAMPLER_DESC samplerDesc = D3D12RootSignature::CreateStaticSampler(
        0,
        D3D12RootSignature::CreateSamplerDesc(
            FilterMode::Linear,
            TextureAddressMode::Clamp,
            D3D12_FLOAT32_MAX
        ),
        ShaderVisibility::Pixel
    );

    D3D12_ROOT_SIGNATURE_DESC rsDesc = D3D12RootSignature::CreateDesc(
        rootParameters, 2,
        &samplerDesc, 1,
        D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT
    );

    if (!gRootSignature.Initialize(device, rsDesc)) {
        Log("[ERROR] Failed to initialize root signature");
        return false;
    }

    D3D12_INPUT_ELEMENT_DESC inputElements[] = {
        D3D12PipelineState::CreateInputElement("POSITION", 0, Format::R32G32B32A32_Float, 0, 0),
        D3D12PipelineState::CreateInputElement("TEXCOORD", 0, Format::R32G32B32A32_Float, 0, 16),
    };

    D3D12_SHADER_BYTECODE emptyGs = {};

    D3D12_GRAPHICS_PIPELINE_STATE_DESC psoDesc = {};
    psoDesc.pRootSignature = gRootSignature.GetRootSignature();
    psoDesc.VS = gVertexShader.GetD3D12Bytecode();
    psoDesc.PS = gPixelShader.GetD3D12Bytecode();
    psoDesc.GS = emptyGs;
    psoDesc.InputLayout.NumElements = 2;
    psoDesc.InputLayout.pInputElementDescs = inputElements;
    psoDesc.RTVFormats[0] = DXGI_FORMAT_R8G8B8A8_UNORM;
    psoDesc.DSVFormat = DXGI_FORMAT_D24_UNORM_S8_UINT;
    psoDesc.SampleDesc.Count = 1;
    psoDesc.SampleDesc.Quality = 0;
    psoDesc.SampleMask = 0xffffffff;
    psoDesc.NumRenderTargets = 1;
    psoDesc.PrimitiveTopologyType = D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE;
    psoDesc.RasterizerState.FillMode = D3D12_FILL_MODE_SOLID;
    psoDesc.RasterizerState.CullMode = D3D12_CULL_MODE_NONE;
    psoDesc.RasterizerState.FrontCounterClockwise = FALSE;
    psoDesc.RasterizerState.DepthClipEnable = TRUE;
    psoDesc.DepthStencilState.DepthEnable = TRUE;
    psoDesc.DepthStencilState.DepthWriteMask = D3D12_DEPTH_WRITE_MASK_ALL;
    psoDesc.DepthStencilState.DepthFunc = D3D12_COMPARISON_FUNC_LESS;
    psoDesc.BlendState.RenderTarget[0].BlendEnable = FALSE;
    psoDesc.BlendState.RenderTarget[0].SrcBlend = D3D12_BLEND_ONE;
    psoDesc.BlendState.RenderTarget[0].DestBlend = D3D12_BLEND_ZERO;
    psoDesc.BlendState.RenderTarget[0].BlendOp = D3D12_BLEND_OP_ADD;
    psoDesc.BlendState.RenderTarget[0].SrcBlendAlpha = D3D12_BLEND_ONE;
    psoDesc.BlendState.RenderTarget[0].DestBlendAlpha = D3D12_BLEND_ZERO;
    psoDesc.BlendState.RenderTarget[0].BlendOpAlpha = D3D12_BLEND_OP_ADD;
    psoDesc.BlendState.RenderTarget[0].RenderTargetWriteMask = D3D12_COLOR_WRITE_ENABLE_ALL;

    if (!gPipelineState.Initialize(device, psoDesc)) {
        Log("[ERROR] Failed to initialize pipeline state");
        return false;
    }

    std::vector<Vertex> vertices;
    std::vector<UINT32> indices;
    GenerateSphere(vertices, indices, 1.0f, 32);
    gIndexCount = (UINT)indices.size();
    Log("[INFO] Generated %d vertices, %d indices", vertices.size(), indices.size());

    if (!gVertexBuffer.InitializeWithData(device, gCommandList.GetCommandList(), vertices.data(), (UINT)(sizeof(Vertex) * vertices.size()), D3D12_RESOURCE_STATE_VERTEX_AND_CONSTANT_BUFFER)) {
        Log("[ERROR] Failed to initialize vertex buffer");
        return false;
    }
    gVertexBuffer.SetStride(sizeof(Vertex));
    gVertexBuffer.SetBufferType(BufferType::Vertex);

    if (!gIndexBuffer.InitializeWithData(device, gCommandList.GetCommandList(), indices.data(), (UINT)(sizeof(UINT32) * indices.size()), D3D12_RESOURCE_STATE_INDEX_BUFFER)) {
        Log("[ERROR] Failed to initialize index buffer");
        return false;
    }
    gIndexBuffer.SetBufferType(BufferType::Index);

    float aspect = 1280.0f / 720.0f;
    gProjectionMatrix = Matrix4x4::Perspective(45.0f * 3.141592f / 180.0f, aspect, 0.1f, 1000.0f);
    gViewMatrix = Matrix4x4::Identity();
    gModelMatrix = Matrix4x4::Translation(Vector3(0.0f, 0.0f, 5.0f));

    float matrices[64];
    Matrix4x4 projTransposed = gProjectionMatrix.Transpose();
    Matrix4x4 viewTransposed = gViewMatrix.Transpose();
    Matrix4x4 modelTransposed = gModelMatrix.Transpose();
    memcpy(matrices, &projTransposed.m[0][0], 64);
    memcpy(matrices + 16, &viewTransposed.m[0][0], 64);
    memcpy(matrices + 32, &modelTransposed.m[0][0], 64);
    memcpy(matrices + 48, &modelTransposed.m[0][0], 64);

    gMVPBuffer.InitializeWithData(device, gCommandList.GetCommandList(), matrices, sizeof(matrices), D3D12_RESOURCE_STATE_VERTEX_AND_CONSTANT_BUFFER);

    if (!LoadTexture("Res/Image/earth.png", gDiffuseTexture, gDiffuseSRV, device, gSRVHeap, gCommandList.GetCommandList(), gCommandAllocator, gCommandQueue)) {
        Log("[ERROR] Failed to load texture");
        return false;
    }

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

void WaitForGPU() {
    gCommandQueue.WaitForIdle();
}

void ExecuteCommandList() {
    gCommandList.Close();
    void* commandLists[] = { gCommandList.GetCommandList() };
    gCommandQueue.ExecuteCommandLists(1, commandLists);
}

void BeginRender() {
    gCurrentRTIndex = gSwapChain.GetCurrentBackBufferIndex();

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

    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);

    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);

    float clearColor[] = { 0.0f, 0.0f, 1.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);
}

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

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

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

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

    if (!RegisterClassEx(&wc)) {
        Log("[ERROR] Failed to register window class");
        return -1;
    }

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

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

    if (!gHWND) {
        Log("[ERROR] Failed to create window");
        return -1;
    }

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

    if (!InitD3D12()) {
        Log("[ERROR] Failed to initialize D3D12");
        return -1;
    }

    RenderDocCapture::Get().SetDevice(gDevice.GetDevice());

    ShowWindow(gHWND, nShowCmd);
    UpdateWindow(gHWND);

    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 {
            if (frameCount > 0) {
                gCommandQueue.WaitForPreviousFrame();
            }

            gCommandAllocator.Reset();
            gCommandList.Reset();

            BeginRender();

            gCommandList.SetRootSignature(gRootSignature.GetRootSignature());
            gCommandList.SetPipelineState(gPipelineState.GetPipelineState());

            ID3D12DescriptorHeap* heaps[] = { gSRVHeap.GetDescriptorHeap() };
            gCommandList.SetDescriptorHeaps(1, heaps);
            gCommandList.SetGraphicsRootConstantBufferView(0, gMVPBuffer.GetResource()->GetGPUVirtualAddress());
            gCommandList.SetGraphicsRootDescriptorTable(1, gSRVHeap.GetGPUDescriptorHandleForHeapStart());

            gCommandList.SetPrimitiveTopology(PrimitiveTopology::TriangleList);
            gCommandList.SetVertexBuffer(0, gVertexBuffer.GetResource(), 0, gVertexBuffer.GetStride());
            gCommandList.SetIndexBuffer(gIndexBuffer.GetResource(), 0, Format::R32_UInt);
            Log("[DEBUG] DrawIndexed with %d indices", gIndexCount);
            gCommandList.DrawIndexed(gIndexCount, 1, 0, 0, 0);

            frameCount++;

            if (frameCount >= targetFrameCount) {
                if (RenderDocCapture::Get().EndCapture()) {
                    Log("[INFO] RenderDoc capture ended");
                }
                Log("[INFO] Reached target frame count %d - taking screenshot!", targetFrameCount);
                WaitForGPU();
                bool screenshotResult = D3D12Screenshot::Capture(
                    gDevice,
                    gCommandQueue,
                    gSwapChain.GetBackBuffer(gCurrentRTIndex),
                    "sphere.ppm"
                );
                if (screenshotResult) {
                    Log("[INFO] Screenshot saved to sphere.ppm");
                } else {
                    Log("[ERROR] Screenshot failed");
                }
                break;
            }

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

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

    gCommandList.Shutdown();
    gCommandAllocator.Shutdown();
    gSwapChain.Shutdown();
    gDevice.Shutdown();

    RenderDocCapture::Get().Shutdown();
    Logger::Get().Shutdown();

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