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

    // Release all pooled textures
    for (auto& pool_pair : m_texture_pool) {
        for (auto& entry : pool_pair.second) {
            if (entry.texture) {
                entry.texture->Release();
                entry.texture = nullptr;
            }
        }
    }
    m_texture_pool.clear();

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

void D3D12Manager::ExecuteCommandListAndWait()
{
    // Close command list
    m_command_list->Close();

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

    // Wait for GPU to finish
    WaitForGPU();

    // Reset command list for next use
    m_command_allocator->Reset();
    m_command_list->Reset(m_command_allocator, nullptr);
}

bool D3D12Manager::SaveTextureToBMP(ID3D12Resource* texture, uint32_t width, uint32_t height, const char* filename)
{
    if (!texture || !filename) {
        printf("[D3D12Manager] Invalid parameters for SaveTextureToBMP\n");
        return false;
    }

    // Readback texture to CPU
    uint8_t* rgba_data = ReadbackTexture(texture, width, height);
    if (!rgba_data) {
        printf("[D3D12Manager] Failed to readback texture\n");
        return false;
    }

    // BMP file header (14 bytes)
    uint32_t file_size = 54 + (width * height * 4);
    uint8_t bmp_header[54] = {
        'B', 'M',  // Signature
        (uint8_t)(file_size), (uint8_t)(file_size >> 8), (uint8_t)(file_size >> 16), (uint8_t)(file_size >> 24),  // File size
        0, 0, 0, 0,  // Reserved
        54, 0, 0, 0,  // Data offset
        40, 0, 0, 0,  // Info header size
        (uint8_t)(width), (uint8_t)(width >> 8), (uint8_t)(width >> 16), (uint8_t)(width >> 24),  // Width
        (uint8_t)(height), (uint8_t)(height >> 8), (uint8_t)(height >> 16), (uint8_t)(height >> 24),  // Height
        1, 0,  // Planes
        32, 0,  // Bits per pixel (32-bit RGBA)
        0, 0, 0, 0,  // Compression (none)
        (uint8_t)(width * height * 4), (uint8_t)((width * height * 4) >> 8),
        (uint8_t)((width * height * 4) >> 16), (uint8_t)((width * height * 4) >> 24),  // Image size
        0, 0, 0, 0,  // X pixels per meter
        0, 0, 0, 0,  // Y pixels per meter
        0, 0, 0, 0,  // Colors used
        0, 0, 0, 0   // Important colors
    };

    // Open file
    FILE* file = nullptr;
    fopen_s(&file, filename, "wb");
    if (!file) {
        printf("[D3D12Manager] Failed to open file: %s\n", filename);
        delete[] rgba_data;
        return false;
    }

    // Write BMP header
    fwrite(bmp_header, 1, 54, file);

    // Convert RGBA to BGRA for BMP and flip vertically
    for (int y = height - 1; y >= 0; y--) {
        for (uint32_t x = 0; x < width; x++) {
            uint32_t src_idx = (y * width + x) * 4;
            uint8_t bgra[4] = {
                rgba_data[src_idx + 2],  // B
                rgba_data[src_idx + 1],  // G
                rgba_data[src_idx + 0],  // R
                rgba_data[src_idx + 3]   // A
            };
            fwrite(bgra, 1, 4, file);
        }
    }

    fclose(file);
    delete[] rgba_data;

    printf("[D3D12Manager] Saved texture to: %s\n", filename);
    return true;
}

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

    printf("[D3D12Manager] NV12 texture created: %ux%u\n", width, height);
    return texture;
}

ID3D12Resource* D3D12Manager::CreateRGBATexture(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_R8G8B8A8_UNORM;
    desc.SampleDesc.Count = 1;
    desc.SampleDesc.Quality = 0;
    desc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;  // Use default layout for RGBA
    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 RGBA texture: 0x%08X\n", hr);
        return nullptr;
    }

    printf("[D3D12Manager] RGBA texture created: %ux%u\n", width, height);
    return texture;
}

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

    // RGBA format has only 1 subresource
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT layout;
    UINT num_rows = 0;
    UINT64 row_size = 0;
    UINT64 total_bytes = 0;

    m_device->GetCopyableFootprints(&desc, 0, 1, 0, &layout, &num_rows, &row_size, &total_bytes);

    // DEBUG: Check if desc dimensions match parameters
    printf("[D3D12Manager::ReadbackTexture] DEBUG:\n");
    printf("  Input params: width=%u, height=%u\n", width, height);
    printf("  desc.Width=%llu, desc.Height=%u\n", desc.Width, desc.Height);
    printf("  layout.Footprint: Width=%u, Height=%u, Depth=%u, RowPitch=%u\n",
           layout.Footprint.Width, layout.Footprint.Height,
           layout.Footprint.Depth, layout.Footprint.RowPitch);
    printf("  num_rows=%u, row_size=%llu, total_bytes=%llu\n", num_rows, row_size, 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 RGBA texture (single subresource)
    D3D12_TEXTURE_COPY_LOCATION src = {};
    src.pResource = texture;
    src.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    src.SubresourceIndex = 0;

    D3D12_TEXTURE_COPY_LOCATION dst = {};
    dst.pResource = readback_buffer;
    dst.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    dst.PlacedFootprint = layout;

    m_command_list->CopyTextureRegion(&dst, 0, 0, 0, &src, 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 for RGBA data
    uint8_t* rgba_data = new uint8_t[width * height * 4];

    // Copy RGBA data (considering row pitch alignment)
    uint8_t* src_ptr = static_cast<uint8_t*>(mapped_data);
    for (uint32_t y = 0; y < height; y++) {
        memcpy(rgba_data + y * width * 4,
               src_ptr + y * layout.Footprint.RowPitch,
               width * 4);
    }

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

    return rgba_data;
}

ID3D12Resource* D3D12Manager::GetOrCreateRGBATexture(uint32_t width, uint32_t height)
{
    TexturePoolKey key = { width, height };

    // Get or create pool for this resolution
    auto& pool = m_texture_pool[key];

    // Try to find an unused texture in the pool
    for (auto& entry : pool) {
        if (!entry.in_use) {
            entry.in_use = true;
            printf("[D3D12Manager] Reusing RGBA texture from pool: %ux%u (pool size: %zu)\n",
                   width, height, pool.size());
            return entry.texture;
        }
    }

    // All textures in pool are in use, check if we can create more
    if (pool.size() < TEXTURE_POOL_SIZE) {
        // Create new texture and add to pool
        ID3D12Resource* texture = CreateRGBATexture(width, height);
        if (texture) {
            TexturePoolEntry entry;
            entry.texture = texture;
            entry.in_use = true;
            pool.push_back(entry);
            printf("[D3D12Manager] Created new RGBA texture for pool: %ux%u (pool size: %zu/%zu)\n",
                   width, height, pool.size(), TEXTURE_POOL_SIZE);
            return texture;
        }
        return nullptr;
    }

    // Pool is full and all textures are in use - this shouldn't happen in normal usage
    printf("[D3D12Manager] WARNING: Texture pool exhausted (%ux%u), all %zu textures in use\n",
           width, height, pool.size());
    return nullptr;
}

void D3D12Manager::ReleaseRGBATexture(ID3D12Resource* texture)
{
    if (!texture) {
        return;
    }

    // Find the texture in the pool and mark it as not in use
    for (auto& pool_pair : m_texture_pool) {
        for (auto& entry : pool_pair.second) {
            if (entry.texture == texture) {
                entry.in_use = false;
                printf("[D3D12Manager] Released RGBA texture back to pool\n");
                return;
            }
        }
    }

    // Texture not found in pool - it might have been created with CreateRGBATexture directly
    printf("[D3D12Manager] WARNING: Released texture not found in pool, releasing directly\n");
    texture->Release();
}