#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/interrupt.h>
#include <linux/irqdomain.h>
#include <linux/msi.h>
#include <linux/pci.h>
#include <linux/eventfd.h>
#include <linux/vfio.h>

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

static irqreturn_t nv_vgpu_msi_handler(int irq, void *data)
{
	struct nv_vgpu_msi_irq *msi = data;
	
	if (msi->trigger)
		eventfd_signal(msi->trigger);
	
	return IRQ_HANDLED;
}

int nv_vgpu_irq_init(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_irq_ctx *ctx;

	if (vgpu->dry_run) {
		nv_vgpu_info(vgpu, "DRY-RUN: IRQ init validation");
		/* Check MSI/MSI-X capability */
		if (pci_find_capability(vgpu->pdev, PCI_CAP_ID_MSIX)) {
			nv_vgpu_info(vgpu, "DRY-RUN: MSI-X capability present");
		} else if (pci_find_capability(vgpu->pdev, PCI_CAP_ID_MSI)) {
			nv_vgpu_info(vgpu, "DRY-RUN: MSI capability present");
		} else {
			nv_vgpu_warn(vgpu, "DRY-RUN: No MSI/MSI-X capability found");
		}
		return 0;
	}

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

	ctx->vgpu = vgpu;
	mutex_init(&ctx->lock);
	ctx->msi_count = 0;
	ctx->msi_enabled = false;
	ctx->msix_enabled = false;

	vgpu->irq_ctx = ctx;
	return 0;
}

void nv_vgpu_irq_cleanup(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;

	if (!ctx)
		return;

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

	nv_vgpu_irq_teardown_msix(vgpu);
	nv_vgpu_irq_teardown_msi(vgpu);

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

int nv_vgpu_irq_setup_msi(struct nv_vgpu_device *vgpu, int num_vectors)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;
	struct pci_dev *pdev = vgpu->pdev;
	int ret, i;

	if (!ctx || num_vectors > NV_VGPU_MAX_MSI_VECTORS)
		return -EINVAL;

	ret = pci_alloc_irq_vectors(pdev, 1, num_vectors, PCI_IRQ_MSI);
	if (ret < 0)
		return ret;

	ctx->msi_count = ret;
	ctx->msi_enabled = true;

	for (i = 0; i < ctx->msi_count; i++) {
		ctx->msis[i].virq = pci_irq_vector(pdev, i);
		ctx->msis[i].vector = i;
		ctx->msis[i].enabled = false;
		ctx->msis[i].trigger = NULL;

		ret = request_irq(ctx->msis[i].virq, nv_vgpu_msi_handler,
				  0, "nv-vgpu-msi", &ctx->msis[i]);
		if (ret) {
			nv_vgpu_warn(vgpu, "Failed to request MSI %d: %d", i, ret);
			ctx->msis[i].virq = 0;
		}
	}

	return 0;
}

void nv_vgpu_irq_teardown_msi(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;
	int i;

	if (!ctx || !ctx->msi_enabled)
		return;

	for (i = 0; i < ctx->msi_count; i++) {
		if (ctx->msis[i].virq)
			free_irq(ctx->msis[i].virq, &ctx->msis[i]);
		if (ctx->msis[i].trigger)
			eventfd_ctx_put(ctx->msis[i].trigger);
	}

	pci_free_irq_vectors(vgpu->pdev);
	ctx->msi_enabled = false;
	ctx->msi_count = 0;
}

int nv_vgpu_irq_setup_msix(struct nv_vgpu_device *vgpu, int num_vectors)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;
	struct pci_dev *pdev = vgpu->pdev;
	int ret, i;

	if (!ctx || num_vectors > NV_VGPU_MAX_MSI_VECTORS)
		return -EINVAL;

	ret = pci_alloc_irq_vectors(pdev, 1, num_vectors, PCI_IRQ_MSIX);
	if (ret < 0)
		return ret;

	ctx->msi_count = ret;
	ctx->msix_enabled = true;

	for (i = 0; i < ctx->msi_count; i++) {
		ctx->msis[i].virq = pci_irq_vector(pdev, i);
		ctx->msis[i].vector = i;
		ctx->msis[i].enabled = false;
		ctx->msis[i].trigger = NULL;

		ret = request_irq(ctx->msis[i].virq, nv_vgpu_msi_handler,
				  0, "nv-vgpu-msix", &ctx->msis[i]);
		if (ret) {
			nv_vgpu_warn(vgpu, "Failed to request MSIX %d: %d", i, ret);
			ctx->msis[i].virq = 0;
		}
	}

	return 0;
}

void nv_vgpu_irq_teardown_msix(struct nv_vgpu_device *vgpu)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;
	int i;

	if (!ctx || !ctx->msix_enabled)
		return;

	for (i = 0; i < ctx->msi_count; i++) {
		if (ctx->msis[i].virq)
			free_irq(ctx->msis[i].virq, &ctx->msis[i]);
		if (ctx->msis[i].trigger)
			eventfd_ctx_put(ctx->msis[i].trigger);
	}

	pci_free_irq_vectors(vgpu->pdev);
	ctx->msix_enabled = false;
	ctx->msi_count = 0;
}

int nv_vgpu_irq_enable(struct nv_vgpu_device *vgpu, int irq_idx)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;

	if (!ctx || irq_idx < 0 || irq_idx >= ctx->msi_count)
		return -EINVAL;

	if (!ctx->msis[irq_idx].virq)
		return -ENODEV;

	ctx->msis[irq_idx].enabled = true;
	enable_irq(ctx->msis[irq_idx].virq);
	return 0;
}

int nv_vgpu_irq_disable(struct nv_vgpu_device *vgpu, int irq_idx)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;

	if (!ctx || irq_idx < 0 || irq_idx >= ctx->msi_count)
		return -EINVAL;

	if (!ctx->msis[irq_idx].virq)
		return -ENODEV;

	disable_irq(ctx->msis[irq_idx].virq);
	ctx->msis[irq_idx].enabled = false;
	return 0;
}

int nv_vgpu_irq_set_trigger(struct nv_vgpu_device *vgpu, int irq_idx, int fd)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;
	struct eventfd_ctx *trigger;

	if (!ctx || irq_idx < 0 || irq_idx >= ctx->msi_count)
		return -EINVAL;

	if (fd < 0) {
		if (ctx->msis[irq_idx].trigger) {
			eventfd_ctx_put(ctx->msis[irq_idx].trigger);
			ctx->msis[irq_idx].trigger = NULL;
		}
		return 0;
	}

	trigger = eventfd_ctx_fdget(fd);
	if (IS_ERR(trigger))
		return PTR_ERR(trigger);

	if (ctx->msis[irq_idx].trigger)
		eventfd_ctx_put(ctx->msis[irq_idx].trigger);

	ctx->msis[irq_idx].trigger = trigger;
	return 0;
}

int nv_vgpu_irq_get_virq(struct nv_vgpu_device *vgpu, int irq_idx)
{
	struct nv_vgpu_irq_ctx *ctx = vgpu->irq_ctx;

	if (!ctx || irq_idx < 0 || irq_idx >= ctx->msi_count)
		return -EINVAL;

	return ctx->msis[irq_idx].virq;
}