/*
 * CHU - Module for driving CH375 in disk mode from Linux 2.0/2.2
 * Copyright (C) 2026 Javier S. Pedro
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 2 of the License, or
 * (at your option) any later version.

 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

/** how many different units we can handle, i.e. /dev/chuda, chudb, etc.*/
#define CHUDISK_MAX_UNITS 1
/** how many bits of the minor number are dedicated to partition number */
#define CHUDISK_MINOR_SHIFT 6
/** how many different partitions we can handle (/dev/chuda1, chuda2, etc.)
 *  which comes from the number of bits available in the minor number. */
#define CHUDISK_MAX_PARTITIONS (1 << CHUDISK_MINOR_SHIFT)
/** total number of minors we can handle. */
#define CHUDISK_MAX_MINORS (CHUDISK_MAX_UNITS << CHUDISK_MINOR_SHIFT)

/* module parameters */
/** configurable major number */
/* by default, reuse one in the experimental space */
static short major = 63;
/** IO base */
static short io = 0x260;
/** skip device proving */
static char force = 0;
/** configurable block size */
static int softblk = 1024;

/* for blk.h: (needs to be defined before blk.h/lnxabs.h is included) */
#define MAJOR_NR major
#define DEVICE_NAME "chud"
#define DEVICE_REQUEST do_chudisk_request
#define DEVICE_NR(device) (MINOR(device)>>CHUDISK_MINOR_SHIFT)
#define DEVICE_ON(device)
#define DEVICE_OFF(device)

#include "lnxabs.h"

#include "ch375lnx.h"
#include <chuver.h>
#include <scsi.h>

#if ENABLE_DLOG
#define dprintf(args...) printk(KERN_DEBUG args)
#else
#define dprintf(args...) /**/
#endif

MODULE_PARM(major, "h");
MODULE_PARM_DESC(major, "major device number");

MODULE_PARM(io, "h");
MODULE_PARM_DESC(io, "io base");

MODULE_PARM(force, "b");
MODULE_PARM_DESC(force, "force loading, skip probing");

MODULE_PARM(softblk, "i");
MODULE_PARM_DESC(softblk, "\"soft\" block size in bytes");

MODULE_DESCRIPTION("CHUDISK CH375 USB disk driver");
MODULE_LICENSE("GPL");

/** hd_structs, used by gendisk to keep track of partition sizes, start points, etc. */
static struct hd_struct chudisk_structs[CHUDISK_MAX_MINORS];
/** used by blk layer to keep track of size of every (minor) blk device. */
static int chudisk_sizes[CHUDISK_MAX_MINORS] = {0};
/** the block size (possibly different for every minor, but we really don't do that) */
static int chudisk_blksizes[CHUDISK_MAX_MINORS] = {0};
/** hardware block size */
static int chudisk_hardsectsizes[CHUDISK_MAX_MINORS] = {0};
/** count open for each unit (not partition) */
static int chudisk_open_count[CHUDISK_MAX_UNITS] = {0};
/** driver data for each unit */
static struct ch375 chudisk_dev[CHUDISK_MAX_UNITS];

/** main gendisk layer data structure, which handles partioning */
static struct gendisk chudisk_gendisk = {
	.major_name = DEVICE_NAME,
	.minor_shift = CHUDISK_MINOR_SHIFT,
	.max_p = CHUDISK_MAX_PARTITIONS,
	.max_nr = CHUDISK_MAX_UNITS,
	.part = chudisk_structs,
	.sizes = chudisk_sizes
};

/** Waits for the CH375 to generate an interrupt, returns result value.
 *  @param timeout max jiffies to wait for an interrupt */
static uint8_t wait_for_ch375_result(struct ch375 *dev, unsigned int timeout)
{
	unsigned long start = jiffies;
	unsigned long expire = start + timeout;

	while (!ch375_poll_interrupt(dev) && time_before(jiffies, expire)) {

	}

	if (ch375_poll_interrupt(dev)) {
		uint8_t status = ch375_get_int_status(dev);
		return status;
	} else {
		printk(KERN_WARNING "chudisk: interrupt timeout! waited %lu jiffies\n", jiffies - start);
		return 0;
	}
}

/** Gets (error) sense data from the USB device (if any)
 *  and prints it to the log. */
static void disk_sense_error(struct ch375 *dev)
{
    char buffer[CH375_BLOCK_SIZE];
    unsigned len;

    ch375_send_command(dev, CH375_DISK_R_SENSE);
    if (wait_for_ch375_result(dev, 1000) != CH375_INT_USB_SUCCESS) {
		printk(KERN_NOTICE "chudisk: could not sense error\n");
        return;
	}

    ch375_send_command(dev, CH375_RD_USB_DATA);
    len = ch375_read_data_block(dev, &buffer[0]);

    printk(KERN_NOTICE "chudisk: sense read data block len = %d\n", len);
    printk(KERN_NOTICE "chudisk: sense packet %hx %hx %hx %hx %hx %hx %hx %hx\n",
            buffer[0], buffer[1], buffer[2], buffer[3],
            buffer[4], buffer[5], buffer[6], buffer[7]);
    printk(KERN_NOTICE "chudisk:              %hx %hx %hx %hx %hx %hx %hx %hx\n",
            buffer[8], buffer[9], buffer[10], buffer[11],
            buffer[12], buffer[13], buffer[14], buffer[15]);
}

/** Read N sectors from the disk
 *  @return number of sectors read
 */
static unsigned disk_read_sectors(struct ch375 *dev, uint8_t *buffer, uint32_t lba, uint8_t num)
{
	uint8_t result;
	unsigned total = 0; // Total bytes written

	dprintf("disk read  buf=%p sector=%u num=%hu\n", buffer, lba, num);

	// Initiate read operation
	ch375_disk_read(dev, lba, num);

	// Now keep going as long as ch375 says to keep going
	while ((result = wait_for_ch375_result(dev, 1000)) == CH375_INT_USB_DISK_READ) {
		unsigned len;

		// Sanity check
		if (total / dev->sec_size >= num) {
			dprintf("read overflow\n");
			break;
		}

		ch375_send_command(dev, CH375_RD_USB_DATA);
		len = ch375_read_data_block(dev, &buffer[total]);

		//dprintf(" disk read data block len = %d\n", len);

		if (len == 0) {
			// Done?
			break;
		}

		total += len;

		// Continue
		ch375_send_command(dev, CH375_DISK_RD_GO);
	}

	dprintf(" disk read  total bytes = %u, sectors = %u, final status = 0x%hx\n",
	        total, total / dev->sec_size, result);

	if (result == CH375_INT_USB_DISK_ERR) {
        disk_sense_error(dev);
    }

	return total / dev->sec_size;
}

/** Write N sectors to the disk
 *  @return number of sectors written
 */
static unsigned disk_write_sectors(struct ch375 *dev, uint8_t *buffer, uint32_t lba, uint8_t num)
{
	uint8_t result;
	unsigned total = 0; // Total bytes written

	dprintf("disk write buf=%p sector=%u num=%hu\n", buffer, lba, num);

	// Initiate operation
	ch375_disk_write(dev, lba, num);

	// Now keep going as long as ch375 says to keep going
	while ((result = wait_for_ch375_result(dev, 1000)) == CH375_INT_USB_DISK_WRITE) {
		const unsigned len = CH375_BLOCK_SIZE;

		// Sanity check
		if (total / dev->sec_size >= num) {
			dprintf("write overflow\n");
			break;
		}

		ch375_send_command(dev, CH375_WR_USB_DATA7);
		ch375_write_data_block(dev, &buffer[total], len);

		total += len;

		// Continue
		ch375_send_command(dev, CH375_DISK_WR_GO);
	}

	dprintf(" disk write total bytes = %u, sectors = %u, final status = 0x%hx\n",
	        total, total / dev->sec_size, result);

    if (result == CH375_INT_USB_DISK_ERR) {
        disk_sense_error(dev);
    }

	return total / dev->sec_size;
}

/** Enables CH375 USB host mode, initializes "disk" mode on it
 *  and connects to a disk.
 *  @return true if connected
  */
static bool chudisk_connect(struct ch375 *dev)
{
	int tries;

	dprintf("switch to usb host\n");
	ch375_set_usb_mode(dev, CH375_USB_MODE_HOST_AUTO_SOF);
	if (wait_for_ch375_result(dev, HZ / 5) != CH375_INT_USB_CONNECT) {
		printk(KERN_INFO "chudisk: could not turn on usb host, likely no device\n");
		return false;
	}

	dprintf("try init disk\n");
	ch375_disk_init(dev);
	if (wait_for_ch375_result(dev, HZ) != CH375_INT_USB_SUCCESS) {
		printk(KERN_INFO "chudisk: could not init disk\n");
		return false;
	}

	// looks like we have to keep polling by asking "is the disk ready?"
	// even if each time we ask we are already polling
	tries = 3;
	while (tries-- > 0) {
		ch375_test_disk_ready(dev);
		if (wait_for_ch375_result(dev, HZ) == CH375_INT_USB_SUCCESS) {
			dprintf("disk now init\n");
			return true;
		}
		mdelay(50);
	}

	printk(KERN_INFO "chudisk: disk failed to init\n");
	return false;
}

/** Tests if there is still a disk connected and initialized. */
static bool is_chudisk_ready(struct ch375 *dev)
{
	ch375_test_disk_ready(dev);
	if (wait_for_ch375_result(dev, HZ/10) == CH375_INT_USB_SUCCESS) {
        return true;
	} else {
		return false;
	}
}

/** Actual main I/O function, called whenever there are pending I/O requests */
static void do_chudisk_request(void)
{
	while (CURRENT) {
		int rc = 0;
		struct hd_struct *hd_struct;
		INIT_REQUEST;

		hd_struct = &chudisk_structs[MINOR(CURRENT->rq_dev)];

		dprintf("do_chudisk_request: cmd %u sector %lu(/%lu) nr_sectors %lu/%lu\n",
		       CURRENT->cmd, CURRENT->sector, hd_struct->nr_sects, CURRENT->current_nr_sectors, CURRENT->nr_sectors);

		// If:
		// the request is for one of our devices
		// and it is in-range of the corresponding partition
		// and respects our max_nr_sectors (256)

		if (CURRENT_DEV < CHUDISK_MAX_UNITS
		        && CURRENT->sector + CURRENT->nr_sectors <= hd_struct->nr_sects
		        && CURRENT->current_nr_sectors <= 0xff) {
			struct ch375 *dev = &chudisk_dev[CURRENT_DEV];
			uint32_t sector = CURRENT->sector + hd_struct->start_sect;
			uint8_t cnt = CURRENT->current_nr_sectors;

			switch (CURRENT->cmd) {
			case READ:
				if (disk_read_sectors(dev, CURRENT->buffer, sector, cnt) == cnt) {
					rc = true; // success
				}
				break;
			case WRITE:
				if (disk_write_sectors(dev, CURRENT->buffer, sector, cnt) == cnt) {
					rc = true; // success
				}
				break;
			default:
				break;
			}
		}

		end_request(rc);
	}
}

/** Called by gendisk to check if removable media may have changed since last time.
 *  If so, gendisk will invalidate buffers and eventually call revalidate.
 *  @return true if media may have changed
 */
 static int chudisk_check_media_change(kdev_t dev)
{
	int unit = DEVICE_NR(dev);
	bool ready;

	dprintf("chudisk_check_media_change: unit=%d\n", unit);

	if (unit >= CHUDISK_MAX_UNITS) {
		return -ENODEV;
	}

	ready = is_chudisk_ready(&chudisk_dev[unit]);

	dprintf("chudisk_check_media_change: unit=%d disk_ready=%d\n", unit, ready);

	// if CH375 is not ready, it means device was unplugged (or never configured to begin with)
	// assume media changed
	return !ready;
}

/** Revalidate means drop all buffers (assuming media changed)
 *  reinitialize all media information (sizes, etc.)
 *  and re-read partition tables. */
static int chudisk_revalidate(kdev_t devno)
{
	int unit = DEVICE_NR(devno);
	struct ch375 *dev = &chudisk_dev[unit];
	uint32_t blk_size;
	uint32_t sector_cnt;
	char buffer[CH375_BLOCK_SIZE];
	int p;

	dprintf("chudisk_revalidate: unit=%d\n", unit);

	if (unit >= CHUDISK_MAX_UNITS) {
		return -ENODEV;
	}

	if (!is_chudisk_ready(dev)) {
		if (!chudisk_connect(dev)) {
			return -ENOMEDIUM;
		}
	}

	ch375_send_command(dev, CH375_DISK_SIZE);
	if (wait_for_ch375_result(dev, HZ / 10) != CH375_INT_USB_SUCCESS) {
        return -EIO;
	}

	ch375_send_command(dev, CH375_RD_USB_DATA);
	if (ch375_read_data_block(dev, buffer) < 8) {
		printk(KERN_WARNING "chudisk: disk size command return less than 8 bytes");
        return -EIO;
	}

    blk_size = ntohl(*(uint32_t*)&buffer[4]);
    sector_cnt = ntohl(*(uint32_t*)&buffer[0]);

	dprintf("chudisk[%d]: sector size %u  num sectors %u\n", unit, blk_size, sector_cnt);

	// TODO Support non-default block_size
	dev->sec_size = blk_size;
	ch375_write_reg(dev, CH375_REG_DISK_PKT_P_SEC, blk_size / CH375_BLOCK_SIZE);

	// Fill in the struct for the raw device, so that we can answer I/O requests on it
	chudisk_structs[unit << CHUDISK_MINOR_SHIFT].nr_sects = sector_cnt;

	// For all minors associated to this unit
	for (p = 0; p < CHUDISK_MAX_PARTITIONS; p++) {
		int minor = (unit << CHUDISK_MINOR_SHIFT) + p;
		kdev_t mdevno = MKDEV(MAJOR_NR, minor);

		fsync_dev(mdevno);
		invalidate_inodes_dev(mdevno);
		invalidate_buffers(mdevno);

		chudisk_hardsectsizes[minor] = blk_size;
	}

	resetup_one_dev(&chudisk_gendisk, unit);

	return 0;
}

/** Fill in guessed hd_geometry from the current media size. */
static int chudisk_get_geom(kdev_t devno, struct hd_geometry *g)
{
	struct hd_struct *hd_struct = &chudisk_structs[DEVICE_NR(devno)];
	struct hd_struct *hd_struct_part = &chudisk_structs[MINOR(devno)];
	unsigned int cylinders, heads, sectors;

	if (hd_struct->nr_sects == 0 || hd_struct_part->nr_sects == 0) {
		return -ENOMEDIUM;
	}

	guess_geometry(hd_struct->nr_sects, &cylinders, &heads, &sectors);
	g->cylinders = cylinders > 0xFFFF ? 0xFFFF : cylinders;
	g->heads = heads;
	g->sectors = sectors;
	g->start = hd_struct_part->start_sect;

	return 0;
}

/** Ioctl handling. */
static int chudisk_ioctl(struct inode *inode, struct file *file, unsigned int cmd, unsigned long arg)
{
	int err;

	dprintf("chudisk_ioctl: unit=%d cmd=%u arg=%lu\n", DEVICE_NR(inode->i_rdev), cmd, arg);

	if (DEVICE_NR(inode->i_rdev) >= CHUDISK_MAX_UNITS) {
		return -ENXIO;
	}

	switch (cmd) {
	case HDIO_GETGEO: {
		struct hd_geometry g;
		if (!arg) return -EINVAL;
		if ((err = chudisk_get_geom(inode->i_rdev, &g))) {
			return err;
		}
		return copy_to_user((struct hd_geometry*)arg, &g, sizeof g) ? -EFAULT : 0;
	}
	case BLKRASET:
		if (!is_sys_admin()) return -EACCES;
		if (arg > 0xff) return -EINVAL;
		read_ahead[MAJOR(inode->i_rdev)] = arg;
		return 0;
	case BLKRAGET:
		if (!arg) return -EINVAL;
		return put_user(read_ahead[MAJOR(inode->i_rdev)], (long *)arg);
	case BLKGETSIZE:
		if (!arg) return -EINVAL;
		return put_user(chudisk_structs[MINOR(inode->i_rdev)].nr_sects, (long*)arg);
	case BLKFLSBUF:
		if (!is_sys_admin()) return -EACCES;
		fsync_dev(inode->i_rdev);
		invalidate_buffers(inode->i_rdev);
		return 0;
	case BLKRRPART:
		if (!is_sys_admin()) return -EACCES;
		return chudisk_revalidate(inode->i_rdev);
	RO_IOCTLS(inode->i_rdev, arg);
	default:
		return -EINVAL;
	}
}

static int chudisk_open(struct inode *inode, struct file *file)
{
	int unit = DEVICE_NR(inode->i_rdev);
	dprintf("chudisk_open: unit=%d\n", unit);
	if (unit >= CHUDISK_MAX_UNITS) {
		return -ENXIO;
	}

	MOD_INC_USE_COUNT;

	if (chudisk_open_count[unit] == 0) {
		// First time open
		dprintf("chudisk: first open\n");
		// Ask gendisk to check if media may have changed (which may call revalidate)
		check_disk_change(inode->i_rdev);
		// if we still have no media, refuse open() call
		if (!is_chudisk_ready(&chudisk_dev[unit])) {
			MOD_DEC_USE_COUNT;
			return -ENOMEDIUM;
		}
	}

	chudisk_open_count[unit]++;

	return 0;
}

static int chudisk_release(struct inode *inode, struct file *file)
{
	int unit = DEVICE_NR(inode->i_rdev);

	dprintf("chudisk_release: unit=%d\n", unit);

	if (unit >= CHUDISK_MAX_UNITS) {
		return -ENXIO;
	}

	chudisk_open_count[unit]--;
	MOD_DEC_USE_COUNT;

	return 0;
}

static struct file_operations blk_fops = {
	.read = block_read,
	.write = block_write,
	.ioctl = chudisk_ioctl,
	.open = chudisk_open,
	.release = chudisk_release,
	.check_media_change = chudisk_check_media_change,
	.revalidate = chudisk_revalidate
};

static __init int chudisk_init(void)
{
	// TODO Here we are currently assuming one unit
	// But it would be trivial to add support for more.
	const int unit = 0;
	struct ch375 *dev = &chudisk_dev[unit];
	int i;

	dev->io_base = io;
	dev->sec_size = 512;

	if (check_region(io, 2)) {
		printk(KERN_WARNING "chudisk: ports at 0x%x not available\n", io);
		return -ENXIO;
	}

	request_region(io, 2, "chudisk");

	if (!force) {
		dprintf("chudisk: resetting CH375 at 0x%hx\n", dev->io_base);
		ch375_reset(dev);
		mdelay(100);
		if (ch375_check_exists(dev, 0xA5) != (uint8_t)~0xA5) {
			printk(KERN_WARNING "chudisk: CH375 not detected at 0x%hx (failed CHECK_EXISTS), change io= param\n", dev->io_base);
			release_region(io, 2);
			return -ENODEV;
		}
	}

	if (register_blkdev(MAJOR_NR, "chud", &blk_fops)) {
		printk(KERN_WARNING "chudisk: unable to get major number %d\n", MAJOR_NR);
		release_region(io, 2);
		return -ENXIO;
	}

	printk(KERN_INFO "chudisk: v%x.%x, found CH375 v.%x at 0x%hx, using major %d\n",
	       CHUDISK_VERSION_MAJOR, CHUDISK_VERSION_MINOR,
	       ch375_get_version(dev), dev->io_base, MAJOR_NR);

	// Set defaults for unit
	for (i = 0; i < CHUDISK_MAX_PARTITIONS; i++) {
		const int minor = (unit << CHUDISK_MINOR_SHIFT) + i;
		chudisk_sizes[minor] = 0;
		chudisk_blksizes[minor] = softblk;
		chudisk_hardsectsizes[minor] = dev->sec_size;
	}

	// Set defaults for everything else
	read_ahead[MAJOR_NR] = 8; // default to 8 sector read-ahead

	blk_dev[MAJOR_NR].request_fn = DEVICE_REQUEST;
	blk_size[MAJOR_NR] = chudisk_sizes;
	blksize_size[MAJOR_NR] = chudisk_blksizes;
	hardsect_size[MAJOR_NR] = chudisk_hardsectsizes;

	// Add ourselves to the genhd linked list
	chudisk_gendisk.major = MAJOR_NR;
	chudisk_gendisk.nr_real = 1;
	chudisk_gendisk.next = gendisk_head;
	gendisk_head = &chudisk_gendisk;

	chudisk_revalidate(MKDEV(MAJOR_NR, unit << CHUDISK_MINOR_SHIFT));

	return 0;
}

static void chudisk_exit(void)
{
	struct ch375 *dev = &chudisk_dev[0];
	struct gendisk **gdp;
	int i;

	// Flush buffers for all minors
	for (i = 0; i < CHUDISK_MAX_MINORS; i++) {
		kdev_t mdevno = MKDEV(MAJOR_NR, i);

		fsync_dev(mdevno);
		invalidate_inodes_dev(mdevno);
		invalidate_buffers(mdevno);
	}

	// Remove from genhd list
	for(gdp = &gendisk_head; *gdp; gdp = &((*gdp)->next)) {
		if (*gdp == &chudisk_gendisk) {
			*gdp = (*gdp)->next;
			break;
		}
	}

	// Remove from blkdev
	blk_dev[MAJOR_NR].request_fn = NULL;
	blk_size[MAJOR_NR] = NULL;
	blksize_size[MAJOR_NR] = NULL;
	hardsect_size[MAJOR_NR] = NULL;
	unregister_blkdev(MAJOR_NR, "chud");

	// Turn off the chip
	ch375_set_usb_mode(dev, CH375_USB_MODE_HOST_DISABLED);
	mdelay(1);
	ch375_enter_sleep(dev);

	// And unregister driver
	release_region(io, 2);

	dprintf("chudisk: bye\n");
}

module_init(chudisk_init);
module_exit(chudisk_exit);