// SPDX-License-Identifier: (GPL-2.0+ OR MIT)
/*
 * Copyright (c) 2025 Amlogic, Inc. All rights reserved.
 */

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/platform_device.h>
#include <linux/of.h>
#include <linux/of_address.h>
#include <linux/of_irq.h>
#include <linux/of_reserved_mem.h>
#include <linux/virtio.h>
#include <linux/virtio_config.h>
#include <linux/virtio_ids.h>
#include <linux/virtio_ring.h>
#include <linux/interrupt.h>
#include <linux/workqueue.h>
#include <linux/io.h>
#include <xen/events.h>
#include <xen/xen.h>

#define DRIVER_NAME "xen_rpmsg"
#define VIRTIO_RPMSG_F_NS 0

/* Number of vrings expected by rpmsg_probe: 0=rx(input), 1=tx(output) */
#define XEN_RPMSG_NUM_VQS 2
#define XEN_RPMSG_VRING_SIZE 256
#define XEN_RPMSG_VRING_ALIGN 0x1000U
#define VRING_SIZE  (0x8000UL)

struct xen_rpmsg_info {
	struct platform_device *pdev;
	struct list_head list;	/* for multi-reg tracking (not currently used) */
	struct virtio_device vdev;

	void __iomem *shm_base;
	resource_size_t shm_size;

	struct virtqueue *vqs[XEN_RPMSG_NUM_VQS];
	int evtchn_irq;
	evtchn_port_t evtchn_port;

	u8 status;
	u64 features;

	struct delayed_work rpmsg_work;
};

#define to_xen_rpmsg_info(_vdev) \
	container_of(_vdev, struct xen_rpmsg_info, vdev)

/* ------------------------------------------------------------------ */
/* virtio_config_ops: get / set                                       */
/* ------------------------------------------------------------------ */

static void xen_rpmsg_get(struct virtio_device *vdev, unsigned int offset,
			  void *buf, unsigned int len)
{
	memset(buf, 0, len);
}

static void xen_rpmsg_set(struct virtio_device *vdev, unsigned int offset,
			  const void *buf, unsigned int len)
{
}

static u32 xen_rpmsg_generation(struct virtio_device *vdev)
{
	return 0;
}

/* ------------------------------------------------------------------ */
/* virtio_config_ops: status                                          */
/* ------------------------------------------------------------------ */

static u8 xen_rpmsg_get_status(struct virtio_device *vdev)
{
	struct xen_rpmsg_info *info = to_xen_rpmsg_info(vdev);

	return info->status;
}

static void xen_rpmsg_set_status(struct virtio_device *vdev, u8 status)
{
	struct xen_rpmsg_info *info = to_xen_rpmsg_info(vdev);

	info->status = status;
}

static void xen_rpmsg_reset(struct virtio_device *vdev)
{
	struct xen_rpmsg_info *info = to_xen_rpmsg_info(vdev);

	info->status = 0;
	info->features = 0;
}

/* ------------------------------------------------------------------ */
/* virtio_config_ops: features                                        */
/* ------------------------------------------------------------------ */

static u64 xen_rpmsg_get_features(struct virtio_device *vdev)
{
	return BIT_ULL(VIRTIO_RPMSG_F_NS);
}

static int xen_rpmsg_finalize_features(struct virtio_device *vdev)
{
	vring_transport_features(vdev);
	return 0;
}

/* ------------------------------------------------------------------ */
/* virtio_config_ops: notify via event channel                        */
/* ------------------------------------------------------------------ */

static bool xen_rpmsg_notify(struct virtqueue *vq)
{
	struct xen_rpmsg_info *info = to_xen_rpmsg_info(vq->vdev);
	/* static atomic_t notify_cnt; */
	/* unsigned int cnt = atomic_inc_return(&notify_cnt); */

	/* dev_info(&info->pdev->dev,
	 *	 "notify #%u: vq[%u] %s evtchn=%u vring_size=%u\n",
	 *	 cnt, vq->index, vq->name ?: "?",
	 *	 info->evtchn_port, virtqueue_get_vring_size(vq));
	 */

	notify_remote_via_evtchn(info->evtchn_port);
	return true;
}

/* ------------------------------------------------------------------ */
/* Event channel interrupt handler                                    */
/* ------------------------------------------------------------------ */

static void xen_rpmsg_work_handler(struct work_struct *work)
{
	struct delayed_work *dwork = to_delayed_work(work);
	struct xen_rpmsg_info *info = container_of(dwork, struct xen_rpmsg_info,
						   rpmsg_work);
	/* static atomic_t work_cnt; */
	/* unsigned int cnt = atomic_inc_return(&work_cnt); */
	int i;

	/* dev_info(&info->pdev->dev,
	 *	 "work #%u: evtchn=%u process %d vqs\n",
	 *	 cnt, info->evtchn_port, XEN_RPMSG_NUM_VQS);
	 */

	for (i = 0; i < XEN_RPMSG_NUM_VQS; i++) {
		if (info->vqs[i])
			vring_interrupt(0, info->vqs[i]);
	}
}

static irqreturn_t xen_rpmsg_irq_handler(int irq, void *dev_id)
{
	struct xen_rpmsg_info *info = dev_id;
	/* static atomic_t irq_cnt; */
	/* unsigned int cnt = atomic_inc_return(&irq_cnt); */

	/* dev_info(&info->pdev->dev,
	 *	 "irq #%u: irq=%d evtchn=%u\n",
	 *	 cnt, irq, info->evtchn_port);
	 */

	schedule_delayed_work(&info->rpmsg_work, 0);

	xen_irq_lateeoi(irq, 0);
	return IRQ_HANDLED;
}

/* ------------------------------------------------------------------ */
/* virtio_config_ops: find_vqs / del_vqs                              */
/* ------------------------------------------------------------------ */

static void xen_rpmsg_del_vqs(struct virtio_device *vdev);

static int xen_rpmsg_find_vqs(struct virtio_device *vdev, unsigned int nvqs,
			      struct virtqueue *vqs[],
			      struct virtqueue_info vqs_info[],
			      struct irq_affinity *desc)
{
	struct xen_rpmsg_info *info = to_xen_rpmsg_info(vdev);
	int i;

	if (nvqs > XEN_RPMSG_NUM_VQS)
		nvqs = XEN_RPMSG_NUM_VQS;

	for (i = 0; i < nvqs; i++) {
		struct virtqueue *vq;
		void *pages;

		/*
		 * Calculate vring offset within SHM:
		 *   vring 0: base + 0
		 *   vring 1: base + vring_size(256, PAGE_SIZE)
		 * This must match the FreeRTOS rpmsg-lite master layout.
		 */
		pages = (void __force *)info->shm_base + i * VRING_SIZE;

		vq = vring_new_virtqueue(i, XEN_RPMSG_VRING_SIZE,
					 XEN_RPMSG_VRING_ALIGN, vdev,
					 false,	/* weak_barriers */
					 false,	/* ctx */
					 pages,
					 xen_rpmsg_notify,
					 vqs_info[i].callback,
					 vqs_info[i].name);
		if (!vq) {
			dev_err(&vdev->dev, "failed to create vq %d\n", i);
			goto error;
		}
		dev_info(&vdev->dev, "vq[%d] %s @ shm+0x%llx (desc=%u)\n",
			 i, vqs_info[i].name,
			 (unsigned long long)(i * VRING_SIZE),
			 XEN_RPMSG_VRING_SIZE);
		vqs[i] = vq;
		info->vqs[i] = vq;
	}

	return 0;

error:
	xen_rpmsg_del_vqs(vdev);
	return -ENOMEM;
}

static void xen_rpmsg_del_vqs(struct virtio_device *vdev)
{
	struct xen_rpmsg_info *info = to_xen_rpmsg_info(vdev);
	struct virtqueue *vq;
	int i;

	for (i = 0; i < XEN_RPMSG_NUM_VQS; i++) {
		vq = info->vqs[i];
		if (vq) {
			vring_del_virtqueue(vq);
			info->vqs[i] = NULL;
		}
	}
}

static const char *xen_rpmsg_bus_name(struct virtio_device *vdev)
{
	return "xen_rpmsg";
}

/* ------------------------------------------------------------------ */
/* virtio_config_ops                                                  */
/* ------------------------------------------------------------------ */

static struct virtio_config_ops xen_rpmsg_config_ops = {
	.get			= xen_rpmsg_get,
	.set			= xen_rpmsg_set,
	.generation		= xen_rpmsg_generation,
	.get_status		= xen_rpmsg_get_status,
	.set_status		= xen_rpmsg_set_status,
	.reset			= xen_rpmsg_reset,
	.find_vqs		= xen_rpmsg_find_vqs,
	.del_vqs		= xen_rpmsg_del_vqs,
	.get_features		= xen_rpmsg_get_features,
	.finalize_features	= xen_rpmsg_finalize_features,
	.bus_name		= xen_rpmsg_bus_name,
};

/* ------------------------------------------------------------------ */
/* Release callback for embedded struct device in virtio_device        */
/* ------------------------------------------------------------------ */

static void xen_rpmsg_device_release(struct device *dev)
{
	struct virtio_device *vdev = container_of(dev, struct virtio_device, dev);
	struct xen_rpmsg_info *info = container_of(vdev, struct xen_rpmsg_info,
						   vdev);

	kfree(info);
}

/* ------------------------------------------------------------------ */
/* Register one virtio device from a reg entry                        */
/* ------------------------------------------------------------------ */

static int xen_rpmsg_register_one(struct platform_device *pdev,
				  struct resource *res, u32 evtchn_port)
{
	struct device *dev = &pdev->dev;
	struct xen_rpmsg_info *info;
	int ret;

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

	info->pdev = pdev;
	info->evtchn_port = evtchn_port;

	/* ---- Map SHM non-cacheable (match FreeRTOS MT_NON_CACHEABLE) ---- */
	info->shm_base = devm_ioremap(dev, res->start, resource_size(res));
	if (IS_ERR(info->shm_base)) {
		dev_err(dev, "ioremap SHM failed: %ld\n",
			PTR_ERR(info->shm_base));
		ret = PTR_ERR(info->shm_base);
		goto free_info;
	}
	info->shm_size = resource_size(res);

	/* ---- Initialize workqueue for deferred vring processing ---- */
	INIT_DELAYED_WORK(&info->rpmsg_work, xen_rpmsg_work_handler);

	/* ---- Bind event channel ---- */
	ret = bind_evtchn_to_irqhandler_lateeoi(evtchn_port,
						xen_rpmsg_irq_handler,
						0, DRIVER_NAME, info);
	if (ret < 0) {
		dev_err(dev, "bind evtchn %u failed: %d\n", evtchn_port, ret);
		goto free_info;
	}
	info->evtchn_irq = ret;

	/* ---- Fill virtio device ---- */
	info->vdev.dev.parent = dev;
	info->vdev.dev.release = xen_rpmsg_device_release;
	info->vdev.id.device = VIRTIO_ID_RPMSG;
	info->vdev.id.vendor = VIRTIO_DEV_ANY_ID;
	info->vdev.config = &xen_rpmsg_config_ops;
	info->vdev.priv = info;

	ret = register_virtio_device(&info->vdev);
	if (ret) {
		dev_err(dev, "register_virtio_device failed: %d\n", ret);
		goto cancel_work;
	}

	dev_info(dev, "xen_rpmsg ready (SHM=%pa size=0x%llx evtchn=%u)\n",
		 &res->start, (unsigned long long)info->shm_size, evtchn_port);
	return 0;

cancel_work:
	cancel_delayed_work_sync(&info->rpmsg_work);
	unbind_from_irqhandler(info->evtchn_irq, info);
	evtchn_put(evtchn_port);
free_info:
	/* shm mapped via devm_ioremap, released with device */
	kfree(info);
	return ret;
}

/* ------------------------------------------------------------------ */
/* Platform driver probe / remove                                     */
/* ------------------------------------------------------------------ */

static int xen_rpmsg_probe(struct platform_device *pdev)
{
	struct device *dev = &pdev->dev;
	u32 evtchn_port;
	int i, ret, nport;

	if (!xen_hvm_domain()) {
		dev_info(dev, "not running in Xen HVM domain, skip probe\n");
		return -ENODEV;
	}

	/*
	 * Init DMA pool from DT memory-region property.
	 * dma_alloc_coherent() in virtio_rpmsg_bus.c will allocate
	 * rpmsg data buffers from this reserved pool.
	 */
	ret = of_reserved_mem_device_init(dev);
	if (ret)
		dev_warn(dev,
			 "no reserved DMA pool, dma_alloc_coherent will fall back to swiotlb\n");

	/* ---- Count xen,evtchn entries (one per rpmsg instance) ---- */
	nport = of_property_count_elems_of_size(dev->of_node, "xen,evtchn",
						sizeof(u32));
	if (nport < 1) {
		dev_err(dev, "missing xen,evtchn property\n");
		ret = -ENODEV;
		goto err;
	}

	/* ---- Loop over each reg entry via of_address_to_resource ---- */
	for (i = 0; ; i++) {
		struct resource res;

		ret = of_address_to_resource(dev->of_node, i, &res);
		if (ret) {
			if (i == 0) {
				dev_err(dev, "missing or invalid reg property\n");
				ret = -ENODEV;
				goto err;
			}
			break;	/* no more regions */
		}

		/* Read per-instance evtchn port */
		ret = of_property_read_u32_index(dev->of_node, "xen,evtchn",
						 i, &evtchn_port);
		if (ret) {
			dev_err(dev, "xen,evtchn[%d] not found\n", i);
			continue;
		}
		dev_info(dev, "evtchn[%d]=%u\n", i, evtchn_port);

		ret = xen_rpmsg_register_one(pdev, &res, evtchn_port);
		if (ret)
			dev_err(dev, "failed to register region %d: %d\n",
				i, ret);
	}

	return 0;

err:
	of_reserved_mem_device_release(dev);
	return ret;
}

static void xen_rpmsg_remove(struct platform_device *pdev)
{
	struct device *dev = &pdev->dev;

	/*
	 * Find and cancel pending work for any registered instance.
	 * Virtio devices are removed via their release callback.
	 */
	of_reserved_mem_device_release(dev);
}

/* ------------------------------------------------------------------ */
/* Device tree match table                                            */
/* ------------------------------------------------------------------ */

static const struct of_device_id xen_rpmsg_of_match[] = {
	{ .compatible = "xen,virtio-rpmsg", },
	{ /* end */ }
};
MODULE_DEVICE_TABLE(of, xen_rpmsg_of_match);

static struct platform_driver xen_rpmsg_platform_driver = {
	.probe	= xen_rpmsg_probe,
	.remove	= xen_rpmsg_remove,
	.driver	= {
		.name		= DRIVER_NAME,
		.of_match_table	= xen_rpmsg_of_match,
	},
};

module_platform_driver(xen_rpmsg_platform_driver);

MODULE_AUTHOR("Xen Dom0less");
MODULE_DESCRIPTION("Xen virtio-rpmsg transport driver");
MODULE_LICENSE("GPL");