#include "utils/dma_alloc.h"

#include <cstdlib>
#include <cerrno>
#include <cstring>
#include <iostream>
#include <mutex>
#include <unordered_map>
#include <vector>

#if defined(__linux__)
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <unistd.h>

// dma_heap ioctl definitions (linux/dma-heap.h may not be available)
struct dma_heap_allocation_data {
    uint64_t len;
    uint32_t fd;
    uint32_t fd_flags;
    uint64_t heap_flags;
};

#define DMA_HEAP_IOC_MAGIC 'H'
#define DMA_HEAP_IOCTL_ALLOC _IOWR(DMA_HEAP_IOC_MAGIC, 0x0, struct dma_heap_allocation_data)

// DMA-BUF sync ioctl
struct dma_buf_sync {
    uint64_t flags;
};

#define DMA_BUF_SYNC_READ (1 << 0)
#define DMA_BUF_SYNC_WRITE (2 << 0)
#define DMA_BUF_SYNC_START (0 << 2)
#define DMA_BUF_SYNC_END (1 << 2)
#define DMA_BUF_BASE 'b'
#define DMA_BUF_IOCTL_SYNC _IOW(DMA_BUF_BASE, 0, struct dma_buf_sync)

#endif

namespace rk3588 {

namespace {

struct DmaBufferPool {
    std::mutex mu;
    std::unordered_map<size_t, std::vector<DmaBuffer*>> free_by_size;
    size_t total_bytes = 0;
    size_t total_buffers = 0;

    bool inited = false;
    bool enabled = true;
    size_t max_bytes = 64ULL * 1024ULL * 1024ULL;
    size_t max_buffers = 16;

    void InitOnce() {
        if (inited) return;
        inited = true;
        if (const char* dis = std::getenv("RK3588_DMA_POOL_DISABLE")) {
            if (std::strcmp(dis, "1") == 0 || std::strcmp(dis, "true") == 0) enabled = false;
        }
        if (const char* mb = std::getenv("RK3588_DMA_POOL_MAX_BYTES")) {
            const long long v = std::atoll(mb);
            if (v > 0) max_bytes = static_cast<size_t>(v);
        }
        if (const char* mc = std::getenv("RK3588_DMA_POOL_MAX_BUFFERS")) {
            const long long v = std::atoll(mc);
            if (v > 0) max_buffers = static_cast<size_t>(v);
        }
    }

    DmaBuffer* Take(size_t size) {
        std::lock_guard<std::mutex> lock(mu);
        InitOnce();
        if (!enabled) return nullptr;
        auto it = free_by_size.find(size);
        if (it == free_by_size.end() || it->second.empty()) return nullptr;
        DmaBuffer* b = it->second.back();
        it->second.pop_back();
        total_bytes -= size;
        total_buffers -= 1;
        return b;
    }

    void Put(DmaBuffer* b) {
        if (!b) return;
        std::lock_guard<std::mutex> lock(mu);
        InitOnce();
        if (!enabled) {
            delete b;
            return;
        }
        const size_t sz = b->size;
        if (sz == 0) {
            delete b;
            return;
        }

        if (total_buffers + 1 > max_buffers || total_bytes + sz > max_bytes) {
            delete b;
            return;
        }

        free_by_size[sz].push_back(b);
        total_bytes += sz;
        total_buffers += 1;
    }
};

static DmaBufferPool& Pool() {
    static DmaBufferPool* p = new DmaBufferPool();
    return *p;
}

}  // namespace

DmaBuffer::~DmaBuffer() {
#if defined(__linux__)
    if (ptr && size > 0) {
        munmap(ptr, size);
        ptr = nullptr;
    }
    if (fd >= 0) {
        close(fd);
        fd = -1;
    }
#endif
    size = 0;
}

DmaBuffer::DmaBuffer(DmaBuffer&& other) noexcept
    : fd(other.fd), ptr(other.ptr), size(other.size) {
    other.fd = -1;
    other.ptr = nullptr;
    other.size = 0;
}

DmaBuffer& DmaBuffer::operator=(DmaBuffer&& other) noexcept {
    if (this != &other) {
#if defined(__linux__)
        if (ptr && size > 0) {
            munmap(ptr, size);
        }
        if (fd >= 0) {
            close(fd);
        }
#endif
        fd = other.fd;
        ptr = other.ptr;
        size = other.size;
        other.fd = -1;
        other.ptr = nullptr;
        other.size = 0;
    }
    return *this;
}

DmaBufferPtr DmaAlloc(size_t alloc_size) {
#if defined(__linux__)
    if (alloc_size == 0) return nullptr;

    if (auto* reused = Pool().Take(alloc_size)) {
        return DmaBufferPtr(reused, [](DmaBuffer* b) { Pool().Put(b); });
    }

    // Try different dma_heap devices.
    // Notes:
    // - Some RGA drivers cannot map buffers whose physical address is >4GB; prefer DMA32 heaps.
    // - CMA heaps are also <4GB friendly but can be small/fragmented under multi-stream load.
    const char* heap_paths[] = {
        // Prefer DMA32 heaps to keep allocations under 4GB.
        // This avoids RGA drivers that cannot map >4GB physical addresses.
        "/dev/dma_heap/system-dma32",
        "/dev/dma_heap/system-uncached-dma32",

        // CMA is also <4GB friendly, but can be small/fragmented under multi-stream load.
        "/dev/dma_heap/cma",
        "/dev/dma_heap/cma-uncached",
        "/dev/dma_heap/linux,cma",

        // Generic system heaps (may allocate >4GB on high-mem systems).
        "/dev/dma_heap/system-uncached",
        "/dev/dma_heap/system",

        // Last resort.
        "/dev/dma_heap/reserved",
        nullptr
    };

    static bool first_success = true;

    int dma_fd = -1;
    const char* used_heap = nullptr;
    int last_errno = 0;
    for (int i = 0; heap_paths[i] != nullptr; ++i) {
        int heap_fd = open(heap_paths[i], O_RDWR | O_CLOEXEC);
        if (heap_fd < 0) {
            last_errno = errno;
            continue;
        }

        dma_heap_allocation_data alloc_data{};
        alloc_data.len = alloc_size;
        alloc_data.fd_flags = O_RDWR | O_CLOEXEC;
        alloc_data.heap_flags = 0;

        if (ioctl(heap_fd, DMA_HEAP_IOCTL_ALLOC, &alloc_data) == 0) {
            dma_fd = static_cast<int>(alloc_data.fd);
            used_heap = heap_paths[i];
            close(heap_fd);
            break;
        }

        last_errno = errno;
        close(heap_fd);
        continue;
    }

    if (dma_fd < 0) {
        std::cerr << "[DmaAlloc] DMA_HEAP_IOCTL_ALLOC failed on all heaps: " << strerror(last_errno) << "\n";
        std::cerr << "[DmaAlloc] tried: ";
        for (int i = 0; heap_paths[i] != nullptr; ++i) {
            std::cerr << heap_paths[i] << " ";
        }
        std::cerr << "\n";
        return nullptr;
    }

    if (first_success) {
        std::cout << "[DmaAlloc] using heap: " << used_heap << "\n";
        first_success = false;
    }

    void* ptr = mmap(nullptr, alloc_size, PROT_READ | PROT_WRITE, MAP_SHARED, dma_fd, 0);
    if (ptr == MAP_FAILED) {
        std::cerr << "[DmaAlloc] mmap failed: " << strerror(errno) << "\n";
        close(dma_fd);
        return nullptr;
    }

    auto* buf = new DmaBuffer();
    buf->fd = dma_fd;
    buf->ptr = ptr;
    buf->size = alloc_size;
    return DmaBufferPtr(buf, [](DmaBuffer* b) { Pool().Put(b); });
#else
    (void)alloc_size;
    std::cerr << "[DmaAlloc] not supported on this platform\n";
    return nullptr;
#endif
}

void DmaSyncStart(DmaBuffer* buf) {
#if defined(__linux__)
    if (!buf || buf->fd < 0) return;
    dma_buf_sync sync{};
    sync.flags = DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ | DMA_BUF_SYNC_WRITE;
    ioctl(buf->fd, DMA_BUF_IOCTL_SYNC, &sync);
#else
    (void)buf;
#endif
}

void DmaSyncEnd(DmaBuffer* buf) {
#if defined(__linux__)
    if (!buf || buf->fd < 0) return;
    dma_buf_sync sync{};
    sync.flags = DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ | DMA_BUF_SYNC_WRITE;
    ioctl(buf->fd, DMA_BUF_IOCTL_SYNC, &sync);
#else
    (void)buf;
#endif
}

void DmaSyncStartFd(int fd) {
#if defined(__linux__)
    if (fd < 0) return;
    dma_buf_sync sync{};
    sync.flags = DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ | DMA_BUF_SYNC_WRITE;
    ioctl(fd, DMA_BUF_IOCTL_SYNC, &sync);
#else
    (void)fd;
#endif
}

void DmaSyncEndFd(int fd) {
#if defined(__linux__)
    if (fd < 0) return;
    dma_buf_sync sync{};
    sync.flags = DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ | DMA_BUF_SYNC_WRITE;
    ioctl(fd, DMA_BUF_IOCTL_SYNC, &sync);
#else
    (void)fd;
#endif
}

}  // namespace rk3588