#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"

int nv_vgpu_memory_init(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_mm_ctx *ctx;
	int ret;

	if (vgpu->dry_run) {
		nv_vgpu_info(vgpu, "DRY-RUN: Memory init validation");
		/* Check IOMMU support */
		if (iommu_present(&pci_bus_type)) {
			nv_vgpu_info(vgpu, "DRY-RUN: IOMMU present");
			ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
			if (!ctx)
				return -ENOMEM;
			ctx->iommu_domain = iommu_domain_alloc(&pci_bus_type);
			if (ctx->iommu_domain) {
				nv_vgpu_info(vgpu, "DRY-RUN: IOMMU domain allocatable");
				iommu_domain_free(ctx->iommu_domain);
			} else {
				nv_vgpu_warn(vgpu, "DRY-RUN: IOMMU domain allocation failed");
			}
			kfree(ctx);
		} else {
			nv_vgpu_warn(vgpu, "DRY-RUN: No IOMMU present");
		}
		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 = iommu_domain_alloc(&pci_bus_type);
	if (!ctx->iommu_domain) {
		ret = -ENODEV;
		goto err_free;
	}

	ret = iommu_attach_device(ctx->iommu_domain, &vgpu->pdev->dev);
	if (ret) {
		nv_vgpu_warn(vgpu, "Failed to attach IOMMU domain: %d", ret);
	}

	vgpu->mm_ctx = ctx;
	return 0;

err_free:
	kfree(ctx);
	return ret;
}

void nv_vgpu_memory_cleanup(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
	struct nv_vgpu_mapping *mapping, *tmp;

	if (!ctx)
		return;

	if (vgpu->dry_run) {
		nv_vgpu_info(vgpu, "DRY-RUN: Memory cleanup (no hardware resources to release)");
		if (ctx->iommu_domain)
			iommu_domain_free(ctx->iommu_domain);
		kfree(ctx);
		vgpu->mm_ctx = NULL;
		return;
	}

	mutex_lock(&ctx->lock);
	list_for_each_entry_safe(mapping, tmp, &ctx->mappings, list) {
		if (mapping->mapped) {
			iommu_unmap(ctx->iommu_domain, mapping->iova, mapping->size);
			mapping->mapped = false;
		}
		if (mapping->pinned) {
			int i;
			for (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);
		ctx->mapping_count--;
	}
	mutex_unlock(&ctx->lock);

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

static struct nv_vgpu_mapping *nv_vgpu_find_mapping(struct nv_vgpu_mm_ctx *ctx, u64 iova, u64 size)
{
	struct nv_vgpu_mapping *mapping;

	list_for_each_entry(mapping, &ctx->mappings, list) {
		if (mapping->iova == iova && mapping->size == size)
			return mapping;
	}
	return 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_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;

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

	mutex_lock(&ctx->lock);

	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;
	list_add_tail(&mapping->list, &ctx->mappings);
	ctx->mapping_count++;

	mutex_unlock(&ctx->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);
	mutex_unlock(&ctx->lock);
	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_mapping *mapping;
	u64 iova = 0;
	int ret = 0;

	if (!ctx)
		return -ENODEV;

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

	return ret;
}

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_mapping *mapping;
	u64 map_iova;
	int ret;

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

	mutex_lock(&ctx->lock);
	mapping = nv_vgpu_find_mapping(ctx, 0, 0);
	if (!mapping || mapping->gpa != gpa || mapping->size != size) {
		list_for_each_entry(mapping, &ctx->mappings, list) {
			if (mapping->gpa == gpa && mapping->size == size)
				break;
		}
		if (&mapping->list == &ctx->mappings) {
			mutex_unlock(&ctx->lock);
			return -ENOENT;
		}
	}

	if (!mapping->pinned) {
		mutex_unlock(&ctx->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(&ctx->lock);
		return -ENOMEM;
	}

	mapping->iova = map_iova;
	mapping->mapped = true;
	*iova = map_iova;

	ret = iommu_map(ctx->iommu_domain, map_iova, gpa, size,
			IOMMU_READ | IOMMU_WRITE | IOMMU_CACHE, GFP_KERNEL);
	if (ret) {
		dma_unmap_page(&vgpu->pdev->dev, map_iova, size, DMA_BIDIRECTIONAL);
		mapping->mapped = false;
		mapping->iova = 0;
	}

	mutex_unlock(&ctx->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_mapping *mapping;
	int ret = 0;

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

	mutex_lock(&ctx->lock);
	mapping = nv_vgpu_find_mapping(ctx, iova, size);
	if (!mapping || !mapping->mapped) {
		mutex_unlock(&ctx->lock);
		return -ENOENT;
	}

	ret = iommu_unmap(ctx->iommu_domain, iova, size);
	if (ret == 0) {
		dma_unmap_page(&vgpu->pdev->dev, iova, size, DMA_BIDIRECTIONAL);
		mapping->mapped = false;
		mapping->iova = 0;
	}

	mutex_unlock(&ctx->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_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;

	mutex_lock(&ctx->lock);
	list_for_each_entry(mapping, &ctx->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(&ctx->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;
}