#pragma once

#include <RHI\IRHIDevice.h>
#include <RHI\IRHIResources.h>
#include <d3d12.h>
#include <dxgi1_4.h>
#include <wrl\client.h>
#include <windows.h>

namespace XCEngine {
namespace RHI {

class D3D12Device;

class D3D12CommandQueue : public ICommandQueue {
public:
    D3D12CommandQueue(D3D12Device* device, ID3D12CommandQueue* queue);
    ~D3D12CommandQueue() override;
    
    void ExecuteCommandLists(void** lists, uint32_t count) override;
    void Signal(void* fence, uint64_t value) override;
    void Wait(void* fence, uint64_t value) override;
    uint64_t GetTimestampFrequency() const override;
    
    ID3D12CommandQueue* GetNativeQueue() const { return m_commandQueue.Get(); }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12CommandQueue> m_commandQueue;
};

class D3D12CommandAllocator : public ICommandAllocator {
public:
    D3D12CommandAllocator(D3D12Device* device, ID3D12CommandAllocator* allocator);
    ~D3D12CommandAllocator() override;
    
    void Reset() override;
    
    ID3D12CommandAllocator* GetNativeAllocator() const { return m_commandAllocator.Get(); }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12CommandAllocator> m_commandAllocator;
};

class D3D12Fence : public IFence {
public:
    D3D12Fence(D3D12Device* device, ID3D12Fence* fence);
    ~D3D12Fence() override;
    
    uint64_t GetCompletedValue() const override;
    void Signal(uint64_t value) override;
    void Wait(uint64_t value) override;
    void Wait(uint64_t value, uint64_t timeoutMs) override;
    
    ID3D12Fence* GetNativeFence() const { return m_fence.Get(); }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12Fence> m_fence;
    HANDLE m_fenceEvent = nullptr;
};

class D3D12DescriptorHeap : public IDescriptorHeap {
public:
    D3D12DescriptorHeap(D3D12Device* device, ID3D12DescriptorHeap* heap, DescriptorHeapType type);
    ~D3D12DescriptorHeap() override;
    
    DescriptorHeapType GetType() const override { return m_type; }
    uint32_t GetDescriptorCount() const override { return m_descriptorCount; }
    void* GetCPUDescriptorHandle(uint32_t index) const override;
    uint64_t GetGPUDescriptorHandle(uint32_t index) const override;
    void SetName(const char* name) override;
    
    ID3D12DescriptorHeap* GetNativeHeap() const { return m_descriptorHeap.Get(); }
    uint32_t GetDescriptorSize() const { return m_descriptorSize; }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12DescriptorHeap> m_descriptorHeap;
    DescriptorHeapType m_type;
    uint32_t m_descriptorCount = 0;
    uint32_t m_descriptorSize = 0;
};

class D3D12SwapChain : public ISwapChain {
public:
    D3D12SwapChain(D3D12Device* device);
    ~D3D12SwapChain() override;
    
    bool Initialize(const SwapChainDesc& desc, void* windowHandle) override;
    void Shutdown() override;
    bool Present() override;
    bool Resize(uint32_t width, uint32_t height) override;
    uint32_t GetCurrentBufferIndex() const override;
    void* GetBuffer(uint32_t index) override;
    void* GetCurrentRenderTarget() override;
    void* GetDepthStencil() override;
    void SetFullscreen(bool fullscreen) override;
    bool IsFullscreen() const override;
    
    IDXGISwapChain3* GetNativeSwapChain() const { return m_swapChain.Get(); }
    ID3D12Resource* GetBufferResource(uint32_t index) const { return m_buffers[index].Get(); }
    D3D12DescriptorHeap* GetRTVHeap() const { return m_rtvHeap; }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<IDXGISwapChain3> m_swapChain;
    Microsoft::WRL::ComPtr<ID3D12Resource> m_buffers[16];
    D3D12DescriptorHeap* m_rtvHeap = nullptr;
    uint32_t m_bufferCount = 0;
    bool m_vsync = false;
    bool m_fullscreen = false;
};

class D3D12RootSignature : public IRootSignature {
public:
    D3D12RootSignature(D3D12Device* device);
    ~D3D12RootSignature() override;
    
    bool Initialize(const RootSignatureDesc& desc) override;
    void SetName(const char* name) override;
    void* GetNativeRootSignature() const override;
    
    ID3D12RootSignature* GetNative() const { return m_rootSignature.Get(); }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12RootSignature> m_rootSignature;
};

class D3D12PipelineState : public IPipelineState {
public:
    D3D12PipelineState(D3D12Device* device);
    ~D3D12PipelineState() override;
    
    bool Initialize(const PipelineDesc& desc);
    void SetName(const char* name) override;
    void* GetNativePipelineState() const override;
    
    ID3D12PipelineState* GetNative() const { return m_pipelineState.Get(); }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12PipelineState> m_pipelineState;
};

class D3D12CommandList : public ICommandList {
public:
    D3D12CommandList(D3D12Device* device, ID3D12GraphicsCommandList* list);
    ~D3D12CommandList() override;
    
    void Reset(void* allocator) override;
    void Close() override;
    void SetPipelineState(IPipelineState* pso) override;
    void SetRootSignature(IRootSignature* signature) override;
    void SetPrimitiveTopology(PrimitiveTopology topology) override;
    void SetVertexBuffer(uint32_t slot, IResource* buffer, uint32_t offset, uint32_t stride) override;
    void SetIndexBuffer(IResource* buffer, uint32_t offset) override;
    void SetDescriptorHeap(IDescriptorHeap* heap) override;
    void SetGraphicsDescriptorTable(uint32_t rootParameterIndex, uint64_t baseDescriptor) override;
    void SetComputeDescriptorTable(uint32_t rootParameterIndex, uint64_t baseDescriptor) override;
    void DrawInstanced(uint32_t vertexCountPerInstance, uint32_t instanceCount, uint32_t startVertex, uint32_t startInstance) override;
    void DrawIndexedInstanced(uint32_t indexCountPerInstance, uint32_t instanceCount, uint32_t startIndex, int32_t baseVertex, uint32_t startInstance) override;
    void Dispatch(uint32_t x, uint32_t y, uint32_t z) override;
    void SetViewports(const Viewport* viewports, uint32_t count) override;
    void SetScissorRects(const Rect* rects, uint32_t count) override;
    void SetRenderTargets(void** targets, uint32_t count, void* depthStencil) override;
    void ClearRenderTargetView(void* target, const float color[4]) override;
    void ClearDepthStencilView(void* depth, float depthValue, uint8_t stencil) override;
    void CopyResource(IResource* dst, const IResource* src) override;
    void CopyBuffer(IResource* dst, uint64_t dstOffset, const IResource* src, uint64_t srcOffset, uint64_t size) override;
    void ResourceBarrier(IResource* resource, ResourceStateFlag before, ResourceStateFlag after) override;
    void* GetNativeCommandList() const override;
    
    ID3D12GraphicsCommandList* GetNative() const { return m_commandList.Get(); }
    D3D12Device* GetDevice() const { return m_device; }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12GraphicsCommandList> m_commandList;
};

class D3D12Resource : public IResource {
public:
    D3D12Resource(D3D12Device* device, ID3D12Resource* resource);
    ~D3D12Resource() override;
    
    void* GetNativeResource() const override { return m_resource.Get(); }
    ID3D12Resource* GetNative() const { return m_resource.Get(); }
    
private:
    D3D12Device* m_device = nullptr;
    Microsoft::WRL::ComPtr<ID3D12Resource> m_resource;
};

class D3D12Device : public IRHIDevice {
public:
    D3D12Device();
    ~D3D12Device() override;
    
    GraphicsAPI GetAPI() const override { return GraphicsAPI::Direct3D12; }
    const char* GetAPIName() const override { return "Direct3D 12"; }
    
    bool Initialize(void* windowHandle, uint32_t width, uint32_t height) override;
    void Shutdown() override;
    
    bool CreateCommandQueue(ICommandQueue** queue, const CommandQueueDesc& desc) override;
    bool CreateCommandAllocator(ICommandAllocator** allocator) override;
    bool CreateCommandList(ICommandList** list, ICommandAllocator* allocator) override;
    bool CreateFence(IFence** fence) override;
    
    bool CreateDescriptorHeap(IDescriptorHeap** heap, const DescriptorHeapDesc& desc) override;
    
    bool CreateRootSignature(IRootSignature** signature, const RootSignatureDesc& desc) override;
    bool CreatePipelineState(IPipelineState** pso, const PipelineDesc& desc) override;
    
    bool CreateSwapChain(ISwapChain** swapChain, const SwapChainDesc& desc, void* windowHandle) override;
    
    ICommandQueue* GetCommandQueue() override { return m_commandQueue; }
    ISwapChain* GetSwapChain() { return m_swapChain; }
    void* GetNativeDevice() const override { return m_device.Get(); }
    void* GetNativeAdapter() const override { return m_adapter.Get(); }
    
    ID3D12Device* GetD3D12Device() const { return m_device.Get(); }
    IDXGIFactory4* GetDXGIFactory() const { return m_dxgiFactory.Get(); }

private:
    bool CreateDXGIDevice(void* windowHandle);
    
    Microsoft::WRL::ComPtr<ID3D12Device> m_device;
    Microsoft::WRL::ComPtr<IDXGIAdapter1> m_adapter;
    Microsoft::WRL::ComPtr<IDXGIFactory4> m_dxgiFactory;
    
    D3D12CommandQueue* m_commandQueue = nullptr;
    D3D12SwapChain* m_swapChain = nullptr;
};

} // namespace RHI
} // namespace XCEngine