#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/vmalloc.h>

#include "nvidia-vgpu-vfio.h"
#include "ioctl.h"
#include "bar_emulation.h"
#include "gsp_comm.h"
#include "memory.h"
#include "irq.h"
#include "migration.h"

int nv_vgpu_ioctl_get_info(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_get_info info;
	int i;

	memset(&info, 0, sizeof(info));

	mutex_lock(&vgpu->lock);
	info.vgpu_type = vgpu->type;
	info.vgpu_id = vgpu->vgpu_id;
	info.gpu_instance_id = vgpu->gpu_instance_id;
	info.compute_instance_id = vgpu->compute_instance_id;
	info.state = vgpu->state;
	info.fb_base = vgpu->fb_base;
	info.fb_size = vgpu->fb_size;
	info.bar0_config_base = vgpu->bar0_config_base;
	info.num_bars = 0;

	for (i = 0; i < NV_VGPU_MAX_REGIONS; i++) {
		if (vgpu->bars[i].size > 0) {
			info.bars[info.num_bars].index = i;
			info.bars[info.num_bars].gpa = vgpu->bars[i].gpa;
			info.bars[info.num_bars].size = vgpu->bars[i].size;
			info.bars[info.num_bars].iova = vgpu->bars[i].iova;
			info.bars[info.num_bars].mapped = vgpu->bars[i].mapped;
			info.bars[info.num_bars].trapped = vgpu->bars[i].trapped;
			info.num_bars++;
		}
	}
	mutex_unlock(&vgpu->lock);

	if (copy_to_user(arg, &info, sizeof(info)))
		return -EFAULT;

	return 0;
}

int nv_vgpu_ioctl_map_bar(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_map_bar params;
	int ret;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	if (params.bar_index >= NV_VGPU_MAX_REGIONS)
		return -EINVAL;

	mutex_lock(&vgpu->lock);

	ret = nv_vgpu_memory_pin(vgpu, params.gpa, params.size);
	if (ret) {
		mutex_unlock(&vgpu->lock);
		return ret;
	}

	ret = nv_vgpu_memory_map(vgpu, params.gpa, params.size, &params.iova);
	if (ret) {
		nv_vgpu_memory_unpin(vgpu, params.gpa, params.size);
		mutex_unlock(&vgpu->lock);
		return ret;
	}

	vgpu->bars[params.bar_index].gpa = params.gpa;
	vgpu->bars[params.bar_index].size = params.size;
	vgpu->bars[params.bar_index].iova = params.iova;
	vgpu->bars[params.bar_index].mapped = true;

	if (params.flags & 0x1) {
		ret = nv_vgpu_bar_trap(vgpu, params.bar_index);
		if (ret)
			nv_vgpu_warn(vgpu, "Failed to trap BAR %d: %d", params.bar_index, ret);
		else
			vgpu->bars[params.bar_index].trapped = true;
	}

	mutex_unlock(&vgpu->lock);

	if (copy_to_user(arg, &params, sizeof(params)))
		return -EFAULT;

	return 0;
}

int nv_vgpu_ioctl_unmap_bar(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_unmap_bar params;
	int ret;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	if (params.bar_index >= NV_VGPU_MAX_REGIONS)
		return -EINVAL;

	mutex_lock(&vgpu->lock);

	if (vgpu->bars[params.bar_index].trapped) {
		nv_vgpu_bar_untrap(vgpu, params.bar_index);
		vgpu->bars[params.bar_index].trapped = false;
	}

	ret = nv_vgpu_memory_unmap(vgpu, params.iova, params.size);
	if (ret) {
		mutex_unlock(&vgpu->lock);
		return ret;
	}

	nv_vgpu_memory_unpin(vgpu, vgpu->bars[params.bar_index].gpa, params.size);

	vgpu->bars[params.bar_index].mapped = false;
	vgpu->bars[params.bar_index].iova = 0;

	mutex_unlock(&vgpu->lock);
	return 0;
}

int nv_vgpu_ioctl_enable_irq(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_irq params;
	int ret;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	ret = nv_vgpu_irq_enable(vgpu, params.irq_index);
	if (ret)
		return ret;

	if (params.fd >= 0) {
		ret = nv_vgpu_irq_set_trigger(vgpu, params.irq_index, params.fd);
		if (ret)
			nv_vgpu_irq_disable(vgpu, params.irq_index);
	}

	return ret;
}

int nv_vgpu_ioctl_disable_irq(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_irq params;
	int ret;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	ret = nv_vgpu_irq_set_trigger(vgpu, params.irq_index, -1);
	if (ret)
		return ret;

	return nv_vgpu_irq_disable(vgpu, params.irq_index);
}

int nv_vgpu_ioctl_gsp_rpc(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_gsp_rpc params;
	void *data = NULL;
	int ret = 0;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	if (params.data_size > 0 && params.data_size <= NV_VGPU_GSP_RPC_MAX_SIZE) {
		data = kmalloc(params.data_size, GFP_KERNEL);
		if (!data)
			return -ENOMEM;

		if (copy_from_user(data, (void __user *)params.data_ptr, params.data_size)) {
			kfree(data);
			return -EFAULT;
		}
	}

	ret = nv_vgpu_gsp_send_rpc(vgpu, data, params.data_size);
	params.result = ret;

	if (data)
		kfree(data);

	if (copy_to_user(arg, &params, sizeof(params)))
		return -EFAULT;

	return 0;
}

int nv_vgpu_ioctl_migrate_save(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_migrate params;
	u8 *data = NULL;
	size_t size = 0;
	int ret;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	ret = nv_vgpu_migration_prepare(vgpu);
	if (ret)
		return ret;

	ret = nv_vgpu_migration_complete(vgpu);
	if (ret)
		return ret;

	data = nv_vgpu_migration_get_data(vgpu, &size);

	if (params.data_ptr && params.data_size >= size) {
		if (copy_to_user((void __user *)params.data_ptr, data, size)) {
			ret = -EFAULT;
		}
	}

	params.data_size = size;
	params.result = ret;

	if (copy_to_user(arg, &params, sizeof(params)))
		return -EFAULT;

	return ret;
}

int nv_vgpu_ioctl_migrate_restore(struct nv_vgpu_device *vgpu, void __user *arg)
{
	struct nv_vgpu_ioctl_migrate params;
	u8 *data = NULL;
	int ret;

	if (copy_from_user(&params, arg, sizeof(params)))
		return -EFAULT;

	if (!params.data_ptr || params.data_size == 0)
		return -EINVAL;

	data = kmalloc(params.data_size, GFP_KERNEL);
	if (!data)
		return -ENOMEM;

	if (copy_from_user(data, (void __user *)params.data_ptr, params.data_size)) {
		kfree(data);
		return -EFAULT;
	}

	ret = nv_vgpu_migration_restore_state(vgpu, data, params.data_size);
	params.result = ret;

	kfree(data);

	if (copy_to_user(arg, &params, sizeof(params)))
		return -EFAULT;

	return ret;
}