#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/list.h>
#include <linux/mm.h>
#include <linux/dma-mapping.h>
#include <linux/iommu.h>
#include <linux/vmalloc.h>
#include <linux/highmem.h>
#include <linux/pfn.h>
#include <linux/scatterlist.h>

#include "nvidia-vgpu-vfio.h"
#include "memory.h"

struct nv_vgpu_simple_mapping {
    u64 gpa;
    u64 size;
    u64 iova;
    struct page **pages;
    int page_count;
    bool pinned;
    bool mapped;
    struct list_head list;
};

static LIST_HEAD(simple_mappings);
static DEFINE_MUTEX(simple_lock);

int nv_vgpu_memory_init(struct nv_vgpu_device *vgpu)
{
    struct nv_vgpu_mm_ctx *ctx;
    
    if (vgpu->dry_run) {
        nv_vgpu_info(vgpu, "DRY-RUN: Memory init validation (simple fallback)");
        return 0;
    }

    ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
    if (!ctx)
        return -ENOMEM;

    ctx->vgpu = vgpu;
    mutex_init(&ctx->lock);
    INIT_LIST_HEAD(&ctx->mappings);
    ctx->mapping_count = 0;
    ctx->iommu_domain = NULL;
    ctx->use_simple = true;

    vgpu->mm_ctx = ctx;
    nv_vgpu_info(vgpu, "Memory init: using simple fallback (no IOMMU)");
    return 0;
}

void nv_vgpu_memory_cleanup(struct nv_vgpu_device *vgpu)
{
    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
    struct nv_vgpu_simple_mapping *mapping, *tmp;
    
    if (!ctx)
        return;

    if (ctx->use_simple) {
        mutex_lock(&simple_lock);
        list_for_each_entry_safe(mapping, tmp, &simple_mappings, list) {
            if (mapping->mapped) {
                dma_unmap_page(&vgpu->pdev->dev, mapping->iova, mapping->size, DMA_BIDIRECTIONAL);
                mapping->mapped = false;
            }
            if (mapping->pinned) {
                for (int i = 0; i < mapping->page_count; i++) {
                    if (mapping->pages[i])
                        put_page(mapping->pages[i]);
                }
                mapping->pinned = false;
            }
            list_del(&mapping->list);
            vfree(mapping->pages);
            kfree(mapping);
        }
        mutex_unlock(&simple_lock);
    } else if (ctx->iommu_domain) {
        iommu_detach_device(ctx->iommu_domain, &vgpu->pdev->dev);
        iommu_domain_free(ctx->iommu_domain);
    }

    kfree(ctx);
    vgpu->mm_ctx = NULL;
}

int nv_vgpu_memory_pin(struct nv_vgpu_device *vgpu, u64 gpa, u64 size)
{
    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
    struct nv_vgpu_simple_mapping *mapping;
    u64 start_pfn = gpa >> PAGE_SHIFT;
    u64 end_pfn = (gpa + size + PAGE_SIZE - 1) >> PAGE_SHIFT;
    int page_count = end_pfn - start_pfn;
    int i, ret = 0;

    if (!ctx)
        return -ENODEV;

    if (!ctx->use_simple)
        return -ENODEV;

    mapping = kzalloc(sizeof(*mapping), GFP_KERNEL);
    if (!mapping)
        return -ENOMEM;

    mapping->gpa = gpa;
    mapping->size = size;
    mapping->page_count = page_count;
    mapping->pages = vzalloc(page_count * sizeof(struct page *));
    if (!mapping->pages) {
        kfree(mapping);
        return -ENOMEM;
    }

    for (i = 0; i < page_count; i++) {
        u64 pfn = start_pfn + i;
        struct page *page = pfn_to_page(pfn);
        if (!page) {
            ret = -EFAULT;
            goto err_unlock;
        }
        get_page(page);
        mapping->pages[i] = page;
    }

    mapping->pinned = true;
    mutex_lock(&simple_lock);
    list_add_tail(&mapping->list, &simple_mappings);
    mutex_unlock(&simple_lock);

    return 0;

err_unlock:
    for (i = 0; i < page_count; i++) {
        if (mapping->pages[i])
            put_page(mapping->pages[i]);
    }
    vfree(mapping->pages);
    kfree(mapping);
    return ret;
}

int nv_vgpu_memory_unpin(struct nv_vgpu_device *vgpu, u64 gpa, u64 size)
{
    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
    struct nv_vgpu_simple_mapping *mapping;
    
    if (!ctx || !ctx->use_simple)
        return -ENODEV;

    mutex_lock(&simple_lock);
    list_for_each_entry(mapping, &simple_mappings, list) {
        if (mapping->gpa == gpa && mapping->size == size) {
            if (mapping->mapped) {
                mutex_unlock(&simple_lock);
                return -EBUSY;
            } else if (mapping->pinned) {
                for (int i = 0; i < mapping->page_count; i++) {
                    if (mapping->pages[i])
                        put_page(mapping->pages[i]);
                }
                mapping->pinned = false;
            }
            list_del(&mapping->list);
            vfree(mapping->pages);
            kfree(mapping);
            break;
        }
    }
    mutex_unlock(&simple_lock);

    return 0;
}

int nv_vgpu_memory_map(struct nv_vgpu_device *vgpu, u64 gpa, u64 size, u64 *iova)
{
    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
    struct nv_vgpu_simple_mapping *mapping;
    u64 map_iova;
    int ret = 0;

    if (!ctx || !ctx->use_simple)
        return -ENODEV;

    mutex_lock(&simple_lock);
    list_for_each_entry(mapping, &simple_mappings, list) {
        if (mapping->gpa == gpa && mapping->size == size)
            break;
    }
    if (&mapping->list == &simple_mappings) {
        mutex_unlock(&simple_lock);
        return -ENOENT;
    }

    if (!mapping->pinned) {
        mutex_unlock(&simple_lock);
        return -EINVAL;
    }

    map_iova = dma_map_page(&vgpu->pdev->dev, mapping->pages[0], 0, size, DMA_BIDIRECTIONAL);
    if (dma_mapping_error(&vgpu->pdev->dev, map_iova)) {
        mutex_unlock(&simple_lock);
        return -ENOMEM;
    }

    mapping->iova = map_iova;
    mapping->mapped = true;
    *iova = map_iova;
    mutex_unlock(&simple_lock);

    return ret;
}

int nv_vgpu_memory_unmap(struct nv_vgpu_device *vgpu, u64 iova, u64 size)
{
    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
    struct nv_vgpu_simple_mapping *mapping;
    int ret = 0;

    if (!ctx || !ctx->use_simple)
        return -ENODEV;

    mutex_lock(&simple_lock);
    list_for_each_entry(mapping, &simple_mappings, list) {
        if (mapping->iova == iova && mapping->size == size) {
            if (!mapping->mapped) {
                mutex_unlock(&simple_lock);
                return -ENOENT;
            }
            dma_unmap_page(&vgpu->pdev->dev, iova, size, DMA_BIDIRECTIONAL);
            mapping->mapped = false;
            mapping->iova = 0;
            break;
        }
    }
    mutex_unlock(&simple_lock);

    return ret;
}

static vm_fault_t nv_vgpu_vm_fault(struct vm_fault *vmf)
{
    struct nv_vgpu_device *vgpu = vmf->vma->vm_private_data;
    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
    struct nv_vgpu_simple_mapping *mapping;
    u64 offset = vmf->address - vmf->vma->vm_start + vmf->vma->vm_pgoff * PAGE_SIZE;
    u64 pfn;
    vm_fault_t ret = VM_FAULT_SIGBUS;

    if (!ctx)
        return VM_FAULT_SIGBUS;

    if (!ctx->use_simple)
        return VM_FAULT_SIGBUS;

    mutex_lock(&simple_lock);
    list_for_each_entry(mapping, &simple_mappings, list) {
        if (mapping->mapped && offset >= mapping->iova &&
            offset < mapping->iova + mapping->size) {
            pfn = (mapping->gpa + offset - mapping->iova) >> PAGE_SHIFT;
            vmf->page = pfn_to_page(pfn);
            get_page(vmf->page);
            ret = VM_FAULT_NOPAGE;
            break;
        }
    }
    mutex_unlock(&simple_lock);

    return ret;
}

static const struct vm_operations_struct nv_vgpu_vm_ops = {
    .fault = nv_vgpu_vm_fault,
};

int nv_vgpu_memory_mmap(struct nv_vgpu_device *vgpu, struct vm_area_struct *vma)
{
    vma->vm_ops = &nv_vgpu_vm_ops;
    vma->vm_private_data = vgpu;
    vm_flags_set(vma, VM_DONTEXPAND | VM_DONTDUMP);
    return 0;
}