#include "D3D12Manager.h"
#include <stdio.h>

D3D12Manager::D3D12Manager()
    : m_device(nullptr)
    , m_command_queue(nullptr)
    , m_command_allocator(nullptr)
    , m_command_list(nullptr)
    , m_fence(nullptr)
    , m_fence_value(0)
    , m_fence_event(nullptr)
{
}

D3D12Manager::~D3D12Manager()
{
    Cleanup();
}

bool D3D12Manager::Initialize()
{
    if (!CreateDevice()) {
        printf("[D3D12Manager] Failed to create D3D12 device\n");
        return false;
    }

    if (!CreateCommandObjects()) {
        printf("[D3D12Manager] Failed to create command objects\n");
        Cleanup();
        return false;
    }

    printf("[D3D12Manager] Initialized successfully\n");
    return true;
}

void D3D12Manager::Cleanup()
{
    WaitForGPU();

    if (m_fence_event) {
        CloseHandle(m_fence_event);
        m_fence_event = nullptr;
    }

    if (m_fence) {
        m_fence->Release();
        m_fence = nullptr;
    }

    if (m_command_list) {
        m_command_list->Release();
        m_command_list = nullptr;
    }

    if (m_command_allocator) {
        m_command_allocator->Release();
        m_command_allocator = nullptr;
    }

    if (m_command_queue) {
        m_command_queue->Release();
        m_command_queue = nullptr;
    }

    if (m_device) {
        m_device->Release();
        m_device = nullptr;
    }
}

bool D3D12Manager::CreateDevice()
{
    // Create DXGI factory
    IDXGIFactory4* factory = nullptr;
    HRESULT hr = CreateDXGIFactory2(0, IID_PPV_ARGS(&factory));
    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create DXGI factory: 0x%08X\n", hr);
        return false;
    }

    // Enumerate adapters and find the first hardware adapter
    IDXGIAdapter1* adapter = nullptr;
    for (UINT i = 0; factory->EnumAdapters1(i, &adapter) != DXGI_ERROR_NOT_FOUND; ++i) {
        DXGI_ADAPTER_DESC1 desc;
        adapter->GetDesc1(&desc);

        // Skip software adapter
        if (desc.Flags & DXGI_ADAPTER_FLAG_SOFTWARE) {
            adapter->Release();
            continue;
        }

        // Try to create D3D12 device
        hr = D3D12CreateDevice(adapter, D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&m_device));
        if (SUCCEEDED(hr)) {
            printf("[D3D12Manager] Created D3D12 device on adapter: %ls\n", desc.Description);
            adapter->Release();
            factory->Release();
            return true;
        }

        adapter->Release();
    }

    factory->Release();
    printf("[D3D12Manager] No suitable D3D12 adapter found\n");
    return false;
}

bool D3D12Manager::CreateCommandObjects()
{
    // Create command queue
    D3D12_COMMAND_QUEUE_DESC queue_desc = {};
    queue_desc.Type = D3D12_COMMAND_LIST_TYPE_DIRECT;
    queue_desc.Flags = D3D12_COMMAND_QUEUE_FLAG_NONE;

    HRESULT hr = m_device->CreateCommandQueue(&queue_desc, IID_PPV_ARGS(&m_command_queue));
    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create command queue: 0x%08X\n", hr);
        return false;
    }

    // Create command allocator
    hr = m_device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT,
        IID_PPV_ARGS(&m_command_allocator));
    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create command allocator: 0x%08X\n", hr);
        return false;
    }

    // Create command list
    hr = m_device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT,
        m_command_allocator, nullptr, IID_PPV_ARGS(&m_command_list));
    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create command list: 0x%08X\n", hr);
        return false;
    }

    // Close command list (we'll reset it when needed)
    m_command_list->Close();

    // Create fence
    hr = m_device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&m_fence));
    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create fence: 0x%08X\n", hr);
        return false;
    }

    m_fence_value = 1;

    // Create fence event
    m_fence_event = CreateEvent(nullptr, FALSE, FALSE, nullptr);
    if (m_fence_event == nullptr) {
        printf("[D3D12Manager] Failed to create fence event\n");
        return false;
    }

    return true;
}

void D3D12Manager::WaitForGPU()
{
    if (!m_command_queue || !m_fence || !m_fence_event) {
        return;
    }

    // Signal the fence
    const UINT64 fence = m_fence_value;
    m_command_queue->Signal(m_fence, fence);
    m_fence_value++;

    // Wait for fence
    if (m_fence->GetCompletedValue() < fence) {
        m_fence->SetEventOnCompletion(fence, m_fence_event);
        WaitForSingleObject(m_fence_event, INFINITE);
    }
}

ID3D12Resource* D3D12Manager::CreateNV12Texture(uint32_t width, uint32_t height)
{
    D3D12_RESOURCE_DESC desc = {};
    desc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
    desc.Width = width;
    desc.Height = height;
    desc.DepthOrArraySize = 1;
    desc.MipLevels = 1;
    desc.Format = DXGI_FORMAT_NV12;
    desc.SampleDesc.Count = 1;
    desc.SampleDesc.Quality = 0;
    desc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;
    desc.Flags = D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS;

    D3D12_HEAP_PROPERTIES heap_props = {};
    heap_props.Type = D3D12_HEAP_TYPE_DEFAULT;

    ID3D12Resource* texture = nullptr;
    HRESULT hr = m_device->CreateCommittedResource(
        &heap_props,
        D3D12_HEAP_FLAG_SHARED,
        &desc,
        D3D12_RESOURCE_STATE_COMMON,
        nullptr,
        IID_PPV_ARGS(&texture));

    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create NV12 texture: 0x%08X\n", hr);
        return nullptr;
    }

    return texture;
}

uint8_t* D3D12Manager::ReadbackTexture(ID3D12Resource* texture, uint32_t width, uint32_t height)
{
    // Get texture description
    D3D12_RESOURCE_DESC desc = texture->GetDesc();

    // Calculate layout for both Y and UV planes
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT layouts[2];
    UINT num_rows[2] = { 0 };
    UINT64 row_sizes[2] = { 0 };
    UINT64 total_bytes = 0;

    m_device->GetCopyableFootprints(&desc, 0, 2, 0, layouts, num_rows, row_sizes, &total_bytes);

    // Create readback buffer
    D3D12_HEAP_PROPERTIES readback_heap_props = {};
    readback_heap_props.Type = D3D12_HEAP_TYPE_READBACK;

    D3D12_RESOURCE_DESC readback_desc = {};
    readback_desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    readback_desc.Width = total_bytes;
    readback_desc.Height = 1;
    readback_desc.DepthOrArraySize = 1;
    readback_desc.MipLevels = 1;
    readback_desc.Format = DXGI_FORMAT_UNKNOWN;
    readback_desc.SampleDesc.Count = 1;
    readback_desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;

    ID3D12Resource* readback_buffer = nullptr;
    HRESULT hr = m_device->CreateCommittedResource(
        &readback_heap_props,
        D3D12_HEAP_FLAG_NONE,
        &readback_desc,
        D3D12_RESOURCE_STATE_COPY_DEST,
        nullptr,
        IID_PPV_ARGS(&readback_buffer));

    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to create readback buffer: 0x%08X\n", hr);
        return nullptr;
    }

    // Reset command list
    m_command_allocator->Reset();
    m_command_list->Reset(m_command_allocator, nullptr);

    // Transition texture to COPY_SOURCE
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Transition.pResource = texture;
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COMMON;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_COPY_SOURCE;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    m_command_list->ResourceBarrier(1, &barrier);

    // Copy Y plane (subresource 0)
    D3D12_TEXTURE_COPY_LOCATION src_y = {};
    src_y.pResource = texture;
    src_y.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    src_y.SubresourceIndex = 0;

    D3D12_TEXTURE_COPY_LOCATION dst_y = {};
    dst_y.pResource = readback_buffer;
    dst_y.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    dst_y.PlacedFootprint = layouts[0];

    m_command_list->CopyTextureRegion(&dst_y, 0, 0, 0, &src_y, nullptr);

    // Copy UV plane (subresource 1)
    D3D12_TEXTURE_COPY_LOCATION src_uv = {};
    src_uv.pResource = texture;
    src_uv.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    src_uv.SubresourceIndex = 1;

    D3D12_TEXTURE_COPY_LOCATION dst_uv = {};
    dst_uv.pResource = readback_buffer;
    dst_uv.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    dst_uv.PlacedFootprint = layouts[1];

    m_command_list->CopyTextureRegion(&dst_uv, 0, 0, 0, &src_uv, nullptr);

    // Transition texture back to COMMON
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_SOURCE;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_COMMON;
    m_command_list->ResourceBarrier(1, &barrier);

    // Execute command list
    m_command_list->Close();
    ID3D12CommandList* cmd_lists[] = { m_command_list };
    m_command_queue->ExecuteCommandLists(1, cmd_lists);

    // Wait for GPU
    WaitForGPU();

    // Map readback buffer
    void* mapped_data = nullptr;
    D3D12_RANGE read_range = { 0, static_cast<SIZE_T>(total_bytes) };
    hr = readback_buffer->Map(0, &read_range, &mapped_data);
    if (FAILED(hr)) {
        printf("[D3D12Manager] Failed to map readback buffer: 0x%08X\n", hr);
        readback_buffer->Release();
        return nullptr;
    }

    // Allocate CPU buffer (NV12 format: Y plane + UV plane)
    uint32_t y_size = width * height;
    uint32_t uv_size = width * (height / 2);
    uint8_t* cpu_buffer = new uint8_t[y_size + uv_size];

    // Copy Y plane
    uint8_t* src_y_data = static_cast<uint8_t*>(mapped_data) + layouts[0].Offset;
    uint8_t* dst_y_data = cpu_buffer;
    for (UINT row = 0; row < height; ++row) {
        memcpy(dst_y_data + row * width, src_y_data + row * layouts[0].Footprint.RowPitch, width);
    }

    // Copy UV plane
    uint8_t* src_uv_data = static_cast<uint8_t*>(mapped_data) + layouts[1].Offset;
    uint8_t* dst_uv_data = cpu_buffer + y_size;
    for (UINT row = 0; row < height / 2; ++row) {
        memcpy(dst_uv_data + row * width, src_uv_data + row * layouts[1].Footprint.RowPitch, width);
    }

    // Unmap and release readback buffer
    readback_buffer->Unmap(0, nullptr);
    readback_buffer->Release();

    return cpu_buffer;
}