#include <windows.h>
#include <filesystem>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <vector>

#include <d3d12.h>
#include <dxgi1_4.h>

#include "XCEngine/RHI/RHIEnums.h"
#include "XCEngine/RHI/RHITypes.h"
#include "XCEngine/RHI/RHIBuffer.h"
#include "XCEngine/RHI/RHIDescriptorPool.h"
#include "XCEngine/RHI/RHIDescriptorSet.h"
#include "XCEngine/RHI/RHIPipelineLayout.h"
#include "XCEngine/RHI/RHIPipelineState.h"
#include "XCEngine/RHI/RHIResourceView.h"
#include "XCEngine/RHI/RHISampler.h"
#include "XCEngine/RHI/RHITexture.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/D3D12SwapChain.h"
#include "XCEngine/RHI/D3D12/D3D12Texture.h"
#include "XCEngine/RHI/D3D12/D3D12ResourceView.h"
#include "XCEngine/RHI/D3D12/D3D12Screenshot.h"
#include "XCEngine/Core/Math/Matrix4.h"
#include "XCEngine/Core/Math/Vector3.h"
#include "XCEngine/Debug/Logger.h"
#include "XCEngine/Debug/ConsoleLogSink.h"
#include "XCEngine/Debug/FileLogSink.h"
#include "XCEngine/Debug/RenderDocCapture.h"
#include "XCEngine/Core/Containers/String.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")

namespace {

struct Vertex {
    float pos[4];
    float uv[2];
};

struct MatrixBufferData {
    Matrix4x4 projection;
    Matrix4x4 view;
    Matrix4x4 model;
};

constexpr float kSphereRadius = 1.0f;
constexpr int kSphereSegments = 32;
constexpr float kPi = 3.14159265358979323846f;

const char kSphereHlsl[] = R"(
Texture2D gDiffuseTexture : register(t0);
SamplerState gSampler : register(s0);

cbuffer MatrixBuffer : register(b0) {
    float4x4 gProjectionMatrix;
    float4x4 gViewMatrix;
    float4x4 gModelMatrix;
};

struct VSInput {
    float4 position : POSITION;
    float2 texcoord : TEXCOORD0;
};

struct PSInput {
    float4 position : SV_POSITION;
    float2 texcoord : TEXCOORD0;
};

PSInput MainVS(VSInput input) {
    PSInput output;
    float4 positionWS = mul(gModelMatrix, input.position);
    float4 positionVS = mul(gViewMatrix, positionWS);
    output.position = mul(gProjectionMatrix, positionVS);
    output.texcoord = input.texcoord;
    return output;
}

float4 MainPS(PSInput input) : SV_TARGET {
    return gDiffuseTexture.Sample(gSampler, input.texcoord);
}
)";

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

std::vector<Vertex> gVertices;
std::vector<uint32_t> gIndices;
RHIBuffer* gVertexBuffer = nullptr;
RHIResourceView* gVertexBufferView = nullptr;
RHIBuffer* gIndexBuffer = nullptr;
RHIResourceView* gIndexBufferView = nullptr;
RHITexture* gTexture = nullptr;
RHIResourceView* gTextureView = nullptr;
RHISampler* gSampler = nullptr;
RHIDescriptorPool* gConstantPool = nullptr;
RHIDescriptorSet* gConstantSet = nullptr;
RHIDescriptorPool* gTexturePool = nullptr;
RHIDescriptorSet* gTextureSet = nullptr;
RHIDescriptorPool* gSamplerPool = nullptr;
RHIDescriptorSet* gSamplerSet = nullptr;
RHIPipelineLayout* gPipelineLayout = nullptr;
RHIPipelineState* gPipelineState = nullptr;

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

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

template <typename T>
void ShutdownAndDelete(T*& object) {
    if (object != nullptr) {
        object->Shutdown();
        delete object;
        object = 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));
}

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* relativePath) {
    return GetExecutableDirectory() / relativePath;
}

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

    segments = segments < 3 ? 3 : segments;

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

        for (int lon = 0; lon <= segments; ++lon) {
            const float theta = (kPi * 2.0f) * static_cast<float>(lon) / static_cast<float>(segments);
            const float sinTheta = sinf(theta);
            const float cosTheta = cosf(theta);

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

            vertex.uv[0] = static_cast<float>(lon) / static_cast<float>(segments);
            vertex.uv[1] = static_cast<float>(lat) / static_cast<float>(segments);
            vertices.push_back(vertex);
        }
    }

    for (int lat = 0; lat < segments; ++lat) {
        for (int lon = 0; lon < segments; ++lon) {
            const uint32_t topLeft = static_cast<uint32_t>(lat * (segments + 1) + lon);
            const uint32_t topRight = topLeft + 1;
            const uint32_t bottomLeft = static_cast<uint32_t>((lat + 1) * (segments + 1) + lon);
            const uint32_t 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);
        }
    }
}

MatrixBufferData CreateMatrixBufferData() {
    const float aspect = 1280.0f / 720.0f;
    const Matrix4x4 projection = Matrix4x4::Perspective(45.0f * 3.141592f / 180.0f, aspect, 0.1f, 1000.0f);
    const Matrix4x4 view = Matrix4x4::Identity();
    const Matrix4x4 model = Matrix4x4::Translation(Vector3(0.0f, 0.0f, 5.0f));

    MatrixBufferData data = {};
    data.projection = projection.Transpose();
    data.view = view.Transpose();
    data.model = model.Transpose();
    return data;
}

bool LoadTexture() {
    const std::filesystem::path texturePath = ResolveRuntimePath("Res/Image/earth.png");
    const std::string texturePathString = texturePath.string();

    stbi_set_flip_vertically_on_load(0);

    int width = 0;
    int height = 0;
    int channels = 0;
    stbi_uc* pixels = stbi_load(texturePathString.c_str(), &width, &height, &channels, STBI_rgb_alpha);
    if (pixels == nullptr) {
        Log("[ERROR] Failed to load texture: %s", texturePathString.c_str());
        return false;
    }

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

    gTexture = gDevice.CreateTexture(
        textureDesc,
        pixels,
        static_cast<size_t>(width) * static_cast<size_t>(height) * 4,
        static_cast<uint32_t>(width) * 4);
    stbi_image_free(pixels);

    if (gTexture == nullptr) {
        Log("[ERROR] Failed to create RHI texture");
        return false;
    }

    ResourceViewDesc textureViewDesc = {};
    textureViewDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    textureViewDesc.dimension = ResourceViewDimension::Texture2D;
    textureViewDesc.mipLevel = 0;
    gTextureView = gDevice.CreateShaderResourceView(gTexture, textureViewDesc);
    if (gTextureView == nullptr) {
        Log("[ERROR] Failed to create texture SRV");
        return false;
    }

    SamplerDesc samplerDesc = {};
    samplerDesc.filter = static_cast<uint32_t>(FilterMode::Linear);
    samplerDesc.addressU = static_cast<uint32_t>(TextureAddressMode::Clamp);
    samplerDesc.addressV = static_cast<uint32_t>(TextureAddressMode::Clamp);
    samplerDesc.addressW = static_cast<uint32_t>(TextureAddressMode::Clamp);
    samplerDesc.mipLodBias = 0.0f;
    samplerDesc.maxAnisotropy = 1;
    samplerDesc.comparisonFunc = static_cast<uint32_t>(ComparisonFunc::Always);
    samplerDesc.borderColorR = 0.0f;
    samplerDesc.borderColorG = 0.0f;
    samplerDesc.borderColorB = 0.0f;
    samplerDesc.borderColorA = 0.0f;
    samplerDesc.minLod = 0.0f;
    samplerDesc.maxLod = 1000.0f;
    gSampler = gDevice.CreateSampler(samplerDesc);
    if (gSampler == nullptr) {
        Log("[ERROR] Failed to create sampler");
        return false;
    }

    DescriptorPoolDesc constantPoolDesc = {};
    constantPoolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    constantPoolDesc.descriptorCount = 1;
    constantPoolDesc.shaderVisible = false;
    gConstantPool = gDevice.CreateDescriptorPool(constantPoolDesc);
    if (gConstantPool == nullptr) {
        Log("[ERROR] Failed to create constant descriptor pool");
        return false;
    }

    DescriptorSetLayoutBinding constantBinding = {};
    constantBinding.binding = 0;
    constantBinding.type = static_cast<uint32_t>(DescriptorType::CBV);
    constantBinding.count = 1;

    DescriptorSetLayoutDesc constantLayoutDesc = {};
    constantLayoutDesc.bindings = &constantBinding;
    constantLayoutDesc.bindingCount = 1;

    gConstantSet = gConstantPool->AllocateSet(constantLayoutDesc);
    if (gConstantSet == nullptr) {
        Log("[ERROR] Failed to allocate constant descriptor set");
        return false;
    }
    const MatrixBufferData matrixData = CreateMatrixBufferData();
    gConstantSet->WriteConstant(0, &matrixData, sizeof(matrixData));

    DescriptorPoolDesc texturePoolDesc = {};
    texturePoolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    texturePoolDesc.descriptorCount = 1;
    texturePoolDesc.shaderVisible = true;
    gTexturePool = gDevice.CreateDescriptorPool(texturePoolDesc);
    if (gTexturePool == nullptr) {
        Log("[ERROR] Failed to create texture descriptor pool");
        return false;
    }

    DescriptorSetLayoutBinding textureBinding = {};
    textureBinding.binding = 0;
    textureBinding.type = static_cast<uint32_t>(DescriptorType::SRV);
    textureBinding.count = 1;

    DescriptorSetLayoutDesc textureLayoutDesc = {};
    textureLayoutDesc.bindings = &textureBinding;
    textureLayoutDesc.bindingCount = 1;

    gTextureSet = gTexturePool->AllocateSet(textureLayoutDesc);
    if (gTextureSet == nullptr) {
        Log("[ERROR] Failed to allocate texture descriptor set");
        return false;
    }
    gTextureSet->Update(0, gTextureView);

    DescriptorPoolDesc samplerPoolDesc = {};
    samplerPoolDesc.type = DescriptorHeapType::Sampler;
    samplerPoolDesc.descriptorCount = 1;
    samplerPoolDesc.shaderVisible = true;
    gSamplerPool = gDevice.CreateDescriptorPool(samplerPoolDesc);
    if (gSamplerPool == nullptr) {
        Log("[ERROR] Failed to create sampler descriptor pool");
        return false;
    }

    DescriptorSetLayoutBinding samplerBinding = {};
    samplerBinding.binding = 0;
    samplerBinding.type = static_cast<uint32_t>(DescriptorType::Sampler);
    samplerBinding.count = 1;

    DescriptorSetLayoutDesc samplerLayoutDesc = {};
    samplerLayoutDesc.bindings = &samplerBinding;
    samplerLayoutDesc.bindingCount = 1;

    gSamplerSet = gSamplerPool->AllocateSet(samplerLayoutDesc);
    if (gSamplerSet == nullptr) {
        Log("[ERROR] Failed to allocate sampler descriptor set");
        return false;
    }
    gSamplerSet->UpdateSampler(0, gSampler);

    return true;
}

GraphicsPipelineDesc CreateSpherePipelineDesc() {
    GraphicsPipelineDesc desc = {};
    desc.pipelineLayout = gPipelineLayout;
    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::D24_UNorm_S8_UInt);
    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 = true;
    desc.depthStencilState.depthWriteEnable = true;
    desc.depthStencilState.depthFunc = static_cast<uint32_t>(ComparisonFunc::Less);
    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 texcoord = {};
    texcoord.semanticName = "TEXCOORD";
    texcoord.semanticIndex = 0;
    texcoord.format = static_cast<uint32_t>(Format::R32G32_Float);
    texcoord.inputSlot = 0;
    texcoord.alignedByteOffset = sizeof(float) * 4;
    desc.inputLayout.elements.push_back(texcoord);

    desc.vertexShader.source.assign(kSphereHlsl, kSphereHlsl + strlen(kSphereHlsl));
    desc.vertexShader.sourceLanguage = ShaderLanguage::HLSL;
    desc.vertexShader.entryPoint = L"MainVS";
    desc.vertexShader.profile = L"vs_5_0";

    desc.fragmentShader.source.assign(kSphereHlsl, kSphereHlsl + strlen(kSphereHlsl));
    desc.fragmentShader.sourceLanguage = ShaderLanguage::HLSL;
    desc.fragmentShader.entryPoint = L"MainPS";
    desc.fragmentShader.profile = L"ps_5_0";

    return desc;
}

bool InitializeSphereResources() {
    GenerateSphere(gVertices, gIndices, kSphereRadius, kSphereSegments);
    if (gVertices.empty() || gIndices.empty()) {
        Log("[ERROR] Failed to generate sphere geometry");
        return false;
    }

    BufferDesc vertexBufferDesc = {};
    vertexBufferDesc.size = static_cast<uint64_t>(gVertices.size() * sizeof(Vertex));
    vertexBufferDesc.stride = sizeof(Vertex);
    vertexBufferDesc.bufferType = static_cast<uint32_t>(BufferType::Vertex);
    gVertexBuffer = gDevice.CreateBuffer(vertexBufferDesc);
    if (gVertexBuffer == nullptr) {
        Log("[ERROR] Failed to create vertex buffer");
        return false;
    }
    gVertexBuffer->SetData(gVertices.data(), gVertices.size() * sizeof(Vertex));
    gVertexBuffer->SetStride(sizeof(Vertex));
    gVertexBuffer->SetBufferType(BufferType::Vertex);

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

    BufferDesc indexBufferDesc = {};
    indexBufferDesc.size = static_cast<uint64_t>(gIndices.size() * sizeof(uint32_t));
    indexBufferDesc.stride = sizeof(uint32_t);
    indexBufferDesc.bufferType = static_cast<uint32_t>(BufferType::Index);
    gIndexBuffer = gDevice.CreateBuffer(indexBufferDesc);
    if (gIndexBuffer == nullptr) {
        Log("[ERROR] Failed to create index buffer");
        return false;
    }
    gIndexBuffer->SetData(gIndices.data(), gIndices.size() * sizeof(uint32_t));
    gIndexBuffer->SetStride(sizeof(uint32_t));
    gIndexBuffer->SetBufferType(BufferType::Index);

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

    if (!LoadTexture()) {
        return false;
    }

    RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.constantBufferCount = 1;
    pipelineLayoutDesc.textureCount = 1;
    pipelineLayoutDesc.samplerCount = 1;
    gPipelineLayout = gDevice.CreatePipelineLayout(pipelineLayoutDesc);
    if (gPipelineLayout == nullptr) {
        Log("[ERROR] Failed to create pipeline layout");
        return false;
    }

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

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

void ShutdownSphereResources() {
    ShutdownAndDelete(gPipelineState);
    ShutdownAndDelete(gPipelineLayout);
    ShutdownAndDelete(gConstantSet);
    ShutdownAndDelete(gTextureSet);
    ShutdownAndDelete(gSamplerSet);
    ShutdownAndDelete(gConstantPool);
    ShutdownAndDelete(gTexturePool);
    ShutdownAndDelete(gSamplerPool);
    ShutdownAndDelete(gSampler);
    ShutdownAndDelete(gTextureView);
    ShutdownAndDelete(gTexture);
    ShutdownAndDelete(gVertexBufferView);
    ShutdownAndDelete(gIndexBufferView);
    ShutdownAndDelete(gVertexBuffer);
    ShutdownAndDelete(gIndexBuffer);
}

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

bool InitD3D12() {
    RHIDeviceDesc deviceDesc = {};
    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;
    }

    if (!gDepthStencil.InitializeDepthStencil(device, gWidth, gHeight)) {
        Log("[ERROR] Failed to initialize depth stencil");
        return false;
    }

    if (!gRTVHeap.Initialize(device, DescriptorHeapType::RTV, 2)) {
        Log("[ERROR] Failed to initialize RTV heap");
        return false;
    }
    gRTVDescriptorSize = gDevice.GetDescriptorHandleIncrementSize(DescriptorHeapType::RTV);

    if (!gDSVHeap.Initialize(device, DescriptorHeapType::DSV, 1)) {
        Log("[ERROR] Failed to initialize DSV heap");
        return false;
    }
    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);

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

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

    if (!InitializeSphereResources()) {
        Log("[ERROR] Failed to initialize sphere resources");
        return false;
    }

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

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

void ExecuteCommandList() {
    gCommandList.Close();
    void* commandLists[] = { &gCommandList };
    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, static_cast<float>(gWidth), static_cast<float>(gHeight), 0.0f, 1.0f };
    Rect scissorRect = { 0, 0, gWidth, gHeight };
    gCommandList.SetViewport(viewport);
    gCommandList.SetScissorRect(scissorRect);

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

void RenderSphere() {
    RHIDescriptorSet* descriptorSets[] = { gConstantSet, gTextureSet, gSamplerSet };
    RHIResourceView* vertexBuffers[] = { gVertexBufferView };
    const uint64_t offsets[] = { 0 };
    const uint32_t strides[] = { sizeof(Vertex) };

    gCommandList.SetPipelineState(gPipelineState);
    gCommandList.SetGraphicsDescriptorSets(0, 3, descriptorSets, gPipelineLayout);
    gCommandList.SetPrimitiveTopology(PrimitiveTopology::TriangleList);
    gCommandList.SetVertexBuffers(0, 1, vertexBuffers, offsets, strides);
    gCommandList.SetIndexBuffer(gIndexBufferView, 0);
    gCommandList.DrawIndexed(static_cast<uint32_t>(gIndices.size()));
}

} // namespace

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

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

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

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

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

    gHWND = CreateWindowExW(
        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();
            RenderSphere();

            frameCount++;

            if (frameCount >= targetFrameCount) {
                Log("[INFO] Reached target frame count %d - taking screenshot!", targetFrameCount);
                ExecuteCommandList();
                if (RenderDocCapture::Get().EndCapture()) {
                    Log("[INFO] RenderDoc capture ended");
                }
                WaitForGPU();
                const 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);
        }
    }

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

    RenderDocCapture::Get().Shutdown();
    Logger::Get().Shutdown();
    return 0;
}