#define NOMINMAX
#include "NanoVDBLoader.h"
#include <windows.h>
#include "BattleFireDirect.h"
#include <nanovdb/io/IO.h>
#include <nanovdb/GridHandle.h>
#include <nanovdb/HostBuffer.h>
#include <cstring>
#include <iostream>

bool LoadNanoVDB(const char* filePath, NanoVDBData& outData, ID3D12GraphicsCommandList* cmdList, ID3D12CommandAllocator* cmdAlloc) {
    try {
        if (cmdAlloc && cmdList) {
            cmdAlloc->Reset();
            cmdList->Reset(cmdAlloc, nullptr);
        }

        nanovdb::GridHandle<nanovdb::HostBuffer> gridHandle = nanovdb::io::readGrid(filePath);
        
        const uint64_t byteSize = gridHandle.buffer().bufferSize();
        const uint64_t elementCount = byteSize / sizeof(uint32_t);
        
        void* bufferData = malloc(byteSize);
        if (!bufferData) {
            std::cerr << "Failed to allocate memory for NanoVDB" << std::endl;
            return false;
        }
        
        memcpy(bufferData, gridHandle.buffer().data(), byteSize);
        
        outData.cpuData = bufferData;
        outData.byteSize = byteSize;
        outData.elementCount = elementCount;
        
        ID3D12Device* device = GetD3DDevice();
        D3D12_HEAP_PROPERTIES heapProps = {};
        heapProps.Type = D3D12_HEAP_TYPE_DEFAULT;
        heapProps.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN;
        heapProps.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN;
        
        D3D12_RESOURCE_DESC resourceDesc = {};
        resourceDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
        resourceDesc.Width = byteSize;
        resourceDesc.Height = 1;
        resourceDesc.DepthOrArraySize = 1;
        resourceDesc.MipLevels = 1;
        resourceDesc.Format = DXGI_FORMAT_UNKNOWN;
        resourceDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
        resourceDesc.SampleDesc.Count = 1;
        
        HRESULT hr = device->CreateCommittedResource(
            &heapProps,
            D3D12_HEAP_FLAG_NONE,
            &resourceDesc,
            D3D12_RESOURCE_STATE_COPY_DEST,
            nullptr,
            IID_PPV_ARGS(&outData.gpuBuffer)
        );
        
        if (FAILED(hr)) {
            std::cerr << "Failed to create GPU buffer for NanoVDB" << std::endl;
            free(bufferData);
            return false;
        }
        
        D3D12_HEAP_PROPERTIES uploadHeapProps = {};
        uploadHeapProps.Type = D3D12_HEAP_TYPE_UPLOAD;
        
        ID3D12Resource* uploadBuffer = nullptr;
        hr = device->CreateCommittedResource(
            &uploadHeapProps,
            D3D12_HEAP_FLAG_NONE,
            &resourceDesc,
            D3D12_RESOURCE_STATE_GENERIC_READ,
            nullptr,
            IID_PPV_ARGS(&uploadBuffer)
        );
        
        if (SUCCEEDED(hr)) {
            void* mappedData = nullptr;
            D3D12_RANGE readRange = {0, 0};
            uploadBuffer->Map(0, &readRange, &mappedData);
            memcpy(mappedData, bufferData, byteSize);
            uploadBuffer->Unmap(0, nullptr);
            
            cmdList->CopyBufferRegion(outData.gpuBuffer, 0, uploadBuffer, 0, byteSize);
            
            D3D12_RESOURCE_BARRIER barrier = {};
            barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
            barrier.Transition.pResource = outData.gpuBuffer;
            barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
            barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_GENERIC_READ;
            cmdList->ResourceBarrier(1, &barrier);
            
            EndCommandList();
            WaitForCompletionOfCommandList();
            uploadBuffer->Release();
            GetCommandAllocator()->Reset();
            cmdList->Reset(GetCommandAllocator(), nullptr);
        }
        
        std::cout << "[NanoVDB] Loaded: " << byteSize << " bytes, " << elementCount << " elements" << std::endl;
        std::cout.flush();
        return true;
    }
    catch (const std::exception& e) {
        std::cerr << "[NanoVDB] Error: " << e.what() << std::endl;
        std::cerr.flush();
        return false;
    }
}

void FreeNanoVDB(NanoVDBData& data) {
    if (data.gpuBuffer) {
        data.gpuBuffer->Release();
        data.gpuBuffer = nullptr;
    }
    if (data.cpuData) {
        free(data.cpuData);
        data.cpuData = nullptr;
    }
    data.byteSize = 0;
    data.elementCount = 0;
}