// [TODO] Clean up headers...
#include <linux/device.h>
#include "asm-generic/errno-base.h"
#include "asm-generic/memory_model.h"
#include "asm/page-def.h"
#include "linux/gfp.h"
#include "linux/mutex.h"
#include "linux/pfn_t.h"
#include "linux/pid.h"
#include <linux/rcupdate.h>
#include <linux/vmalloc.h>
#include <linux/list.h>
#include <linux/types.h>
#include <linux/errno.h>
#include <linux/mm.h>
#include <linux/mm_types.h>
#include <linux/fs.h>
#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/version.h>
#include <linux/ktime.h>

#include <asm/cacheflush_extra.h>

MODULE_AUTHOR("Zk.");
MODULE_DESCRIPTION("4.2.W1: mmap for point of coherency");
MODULE_LICENSE("GPL");

struct my_shmem_page {
	struct page *page;
	struct list_head list;
};
static DEFINE_MUTEX(my_shmem_pages_mtx);

/* [!] READ/WRITE UNDER LOCK */
static LIST_HEAD(my_shmem_pages);

/* [!] READ/WRITE UNDER LOCK */
// static size_t my_shmem_page_count = 0;

static int major;
static struct class* class;
static struct device* dev;
const char* DEV_NAME = "my_shmem";

/* Virtual Memory Area Operations...
 * ============================================================================
 */

/* [TODO]
 * This don't work, not bc. __dcache_clean_poc don't work, but bc. the cache for
 * the *ENTIRE MM* has already been flushed prior to calling this.
 */
static void my_shmem_vmops_close(struct vm_area_struct *vma)
{
	size_t nr_pages_in_cache = list_count_nodes(&my_shmem_pages); 
	size_t nr_pages_of_vma = (vma->vm_end - vma->vm_start) >> PAGE_SHIFT;
	pr_info(
		"[%s] Entered. vma size: %ld; cached pages: %ld.\n", 
		__func__, nr_pages_of_vma, nr_pages_in_cache
	);

	size_t nr_pages_offset = vma->vm_pgoff;
	struct my_shmem_page *entry;
	// u64 clean_time_bgn, clean_time_end;
	u64 runtime;

	mutex_lock(&my_shmem_pages_mtx);
	list_for_each_entry(entry, &my_shmem_pages, list) {
		/* Loop until first page out of offset */
		if (nr_pages_offset) {
			nr_pages_offset--;
			continue;
		}

		/* Early exit when all of vma is done */
		if (!nr_pages_of_vma)
			break;

		/* Inside the window of mapped pages -- flush them up */
		struct page *pg = entry->page;
		ulong kvaddr_bgn = (ulong) page_address(pg);
		ulong kvaddr_end = kvaddr_bgn + PAGE_SIZE;


		pr_info("[%s] Before flush: 0x%px has 0x%lx [+%ld].\n",
			__func__, (void *) kvaddr_bgn, *(ulong *) kvaddr_bgn,
			sizeof(ulong));
		__dcache_clean_poc_dbg(kvaddr_bgn, kvaddr_end, &runtime);
		pr_info("[%s] After flush: 0x%px has 0x%lx [+%ld]. Runtime: %lldns.\n",
			__func__, (void *) kvaddr_bgn, *(ulong *) kvaddr_bgn,
			sizeof(ulong), runtime);

		put_page(pg);

		nr_pages_of_vma--;
	}
	mutex_unlock(&my_shmem_pages_mtx);

	pr_info("[%s] Flushed dcache.\n", __func__);
}

static vm_fault_t my_shmem_vmops_fault(struct vm_fault *vmf)
{
	pr_info("[%s] vm_fault @ 0x%lx (real address 0x%lx).\n",
		__func__, vmf->address, vmf->real_address);

	vm_fault_t ret = 0;
	struct vm_area_struct *vma_of_vmf = vmf->vma;
	ulong vma_pg_offset = vma_of_vmf->vm_pgoff;
	ulong vmf_pg_offset_from_vma_bgn =
		(vmf->address - vma_of_vmf->vm_start) >> PAGE_SHIFT;
	ulong pg_offset = vma_pg_offset + vmf_pg_offset_from_vma_bgn;
	phys_addr_t _dbg_phys_of_page;
	struct page *last_pg;

	mutex_lock(&my_shmem_pages_mtx);
	size_t my_shmem_page_count = list_count_nodes(&my_shmem_pages);
	BUG_ON(pg_offset < my_shmem_page_count);

	/* Allocate the new page(s) */
	ulong nr_pages_to_alloc = pg_offset - my_shmem_page_count + 1;
	pr_info("[%s] Page count %ld, offset %ld -- allocating %ld more...\n",
		__func__, my_shmem_page_count, vmf->pgoff, nr_pages_to_alloc);
	for (; nr_pages_to_alloc > 0; nr_pages_to_alloc--)
	{
		// Allocate page handle in kernel
		struct my_shmem_page *new_page = kzalloc(
			sizeof(struct my_shmem_page), GFP_KERNEL);
		if (!new_page) {
			mutex_unlock(&my_shmem_pages_mtx);
			goto err_ret_no_kmem;
		}

		// Allocate kernel virtual page
		struct page *curr_pg = alloc_page(GFP_USER);
		if (!curr_pg) {
			mutex_unlock(&my_shmem_pages_mtx);
			goto err_ret_no_page;
		}
		pr_info("[%s] Allocated pfn: %ld, kernel vaddr: %px.\n",
			__func__, page_to_pfn(curr_pg), page_to_virt(curr_pg));
		get_page(curr_pg); // For base page refcount

		// Populate page handle
		new_page->page = curr_pg;

		// Add page handle to list
		list_add(&new_page->list, &my_shmem_pages);

		// Fill in last_pg for final return from page fault handler
		last_pg = curr_pg;
	}
	// Fill in vmf's page for return
	get_page(last_pg);
	vmf->page = last_pg;
	// ret = vmf_insert_page(vma_of_vmf, vmf->address, last_pg);
	// [!] YOU DON'T NEED TO CALL REMAP_PFN_RANGE OR FAMILY HERE!!!
	// `__do_fault` allocates PTE prior to calling `vm_ops->fault`,
	// and at return `finish_fault` inserts PTE for given page.
	// [?] I think `vmf_insert_page` etc. are used to insert device pages...
	// They require VM_MIXEDMAP, which in this case don't need (and in fact causes a BUG here.)
	// I think this is also the reason why `remap_pfn_range` might work here,
	// exactly because it does NOT try to add VM_MIXEDMAP, etc.
	_dbg_phys_of_page = page_to_phys(last_pg);

	mutex_unlock(&my_shmem_pages_mtx);
	goto ok_ret_allocated;

err_ret_no_kmem:
	pr_err("[%s] Cannot allocate `struct my_shmem_page` in kernel memory.\n",
		__func__);
	return VM_FAULT_OOM;
err_ret_no_page:
	pr_err("[%s] Cannot allocate requested page for virtual memory.\n",
		__func__);
	return VM_FAULT_OOM;
ok_ret_allocated:
	return ret;
}

static const struct vm_operations_struct my_shmem_vmops = {
	.close = my_shmem_vmops_close,
	.fault = my_shmem_vmops_fault,
};

/* File Operations...
 * ============================================================================
 */
// static int my_shmem_fops_open(struct inode *inode, struct file *filp);

static const struct file_operations my_shmem_fops;

static int my_shmem_fops_mmap(struct file *filp, struct vm_area_struct *vma)
{
	int ret = 0;
	ulong uvaddr = vma->vm_start;
	ulong upgoff = vma->vm_pgoff;
	struct my_shmem_page *curr;
	struct page *pg;
	const unsigned char *fp_name;

	vma->vm_ops = &my_shmem_vmops;

	/* Remap as much as possible */
	mutex_lock(&my_shmem_pages_mtx);
	size_t my_shmem_page_count = list_count_nodes(&my_shmem_pages);
	if (!my_shmem_page_count) {
		mutex_unlock(&my_shmem_pages_mtx);
		goto cleanup_ok;
	}
	list_for_each_entry(curr, &my_shmem_pages, list) {
		/* If userspace virt addr >= vm_end, exit. */
		if (uvaddr >= vma->vm_end)
			break;

		/* Wait until the vm_pgoff-th page, if exists. */
		if (upgoff != 0) {
			upgoff--;
			continue;
		}

		/* Exists allocable page, remap */
		pg = curr->page;
		get_page(pg);
		ret = remap_pfn_range(
			vma, uvaddr, page_to_pfn(pg), PAGE_SIZE,
			vma->vm_page_prot);
		if (ret) {
			mutex_unlock(&my_shmem_pages_mtx);
			goto cleanup_err_remap_pfn_failed;
		}
		pr_info("[%s] Remapped pfn %ld (kvaddr: 0x%px) -> uvaddr: 0x%px.\n",
			__func__, page_to_pfn(pg), page_to_virt(pg), (void *) uvaddr);
		uvaddr += PAGE_SIZE;
	}
	/* May still have unmapped pages, we allocate lazily at fault time. */
	mutex_unlock(&my_shmem_pages_mtx);
	goto cleanup_ok;

cleanup_err_remap_pfn_failed:
	pr_err("[%s] Cannot remap pfn %ld (kvaddr: 0x%px) -> uvaddr: 0x%px: %d.\n",
		__func__, page_to_pfn(pg), page_to_virt(pg), (void *) uvaddr, ret);
	put_page(pg);
	goto exit;
cleanup_ok:
	fp_name = file_dentry(filp)->d_name.name;
	if (uvaddr == vma->vm_end)
		pr_info("[%s] Device file '%s' mmapped for vma: [0x%lx - 0x%lx).\n",
			__func__, fp_name, vma->vm_start, vma->vm_end);
	else
		pr_info("[%s] Device file '%s' mmapped for vma: [0x%lx - 0x%lx..0x%lx).\n",
			__func__, fp_name, vma->vm_start, uvaddr, vma->vm_end);
exit:
	return ret;
}

static int my_shmem_fops_open(struct inode *inode, struct file *filp)
{
	filp->f_op = &my_shmem_fops;

	pr_info("[%s] Device file '%s' opened.\n",
		__func__, file_dentry(filp)->d_name.name);
	return 0;
}

static int my_shmem_fops_release(struct inode *inode, struct file *filp)
{
	pr_info("[%s] Device file '%s' released.\n",
		__func__, file_dentry(filp)->d_name.name);

	/* Garbage collection requires knowing who references which page...
	 * Ideally this is stored in filp->private_data but ah well, oversight.
	 */
	return 0;
}

// static int my_shmem_fops_ioctl(struct file *filp, uint cmd, ulong arg);

// static int my_shmem_fops_fsync(
// 	struct file *filp, loff_t bgn_off, loff_t end_off, int datasync)
// {
// 	pr_info("[%s] Entered.\n", __func__);

// }

static const struct file_operations my_shmem_fops = {
	.owner = THIS_MODULE,
	.open = my_shmem_fops_open,
	.mmap = my_shmem_fops_mmap,
	// .unlocked_ioctl = my_shmem_fops_ioctl,
	// .compat_ioctl = compat_ptr_ioctl,
	.release = my_shmem_fops_release,
};

/* Module init & exit...
 * ============================================================================
 */
static int __init my_shmem_init(void)
{
	/* Register cdev */
	major = register_chrdev(
		0, DEV_NAME, &my_shmem_fops);
	if (major < 0)
		goto err_ret_cdev_reg_failed;

	/* Create device class */
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(6, 4, 0)
	class = class_create(DEV_NAME);
	#else
	// Re: commit dcfbb67e48a2becfce7990386e985b9c45098ee5 (I think?)
	class = class_create(THIS_MODULE, DEV_NAME);
	#endif
	if (IS_ERR(class)) {
		unregister_chrdev(major, DEV_NAME);
		goto err_ret_class_crea_failed;
	}

	/* Create one device */
	dev_t dev_nr = MKDEV(major, 0);
	dev = device_create(
		class, NULL, dev_nr, NULL, DEV_NAME);
	if (IS_ERR(dev)) {
		class_destroy(class);
		unregister_chrdev(major, DEV_NAME);
		goto err_ret_dev_crea_failed;
	}

	pr_info("[%s] Device `%s` built successfully.\n",
		__func__, dev->init_name);
	return 0;

err_ret_cdev_reg_failed:
	pr_err("[%s] Cannot register character dev -- error code %d.\n",
		__func__, major);
	return major;
err_ret_class_crea_failed:
	pr_err("[%s] Cannot create device class -- error code %ld.\n",
		__func__, PTR_ERR(class));
	return (int) PTR_ERR(class);
err_ret_dev_crea_failed:
	pr_err("[%s] Cannot create device -- error code %ld.\n",
		__func__, PTR_ERR(dev));
	return (int) PTR_ERR(dev);
}

static void __exit my_shmem_exit(void)
{
	/* Destroy device */
	device_destroy(class, dev->devt);
	class_destroy(class);
	unregister_chrdev(major, DEV_NAME);
	pr_info("[%s] Device destroyed.\n", __func__);

	/* Free all pages -- I'm not compacting in runtime!!! */
	struct my_shmem_page *page_entry, *tmp;
	mutex_lock(&my_shmem_pages_mtx);
	list_for_each_entry_safe(page_entry, tmp, &my_shmem_pages, list) {
		put_page(page_entry->page);
		free_page((ulong) page_to_virt(page_entry->page));

		// my_shmem_page_count--;
		list_del(&page_entry->list);
		kfree(page_entry);
	}
	mutex_unlock(&my_shmem_pages_mtx);
}

module_init(my_shmem_init);
module_exit(my_shmem_exit);
