#include <socks/printk.h>
#include <socks/pci.h>
#include <socks/device.h>
#include <socks/kext.h>
#include <socks/pci.h>
#include <socks/machine/init.h>
#include <socks/libc/stdio.h>
#include <socks/libc/ctype.h>
#include <arch/irq.h>
#include "ahci.h"

struct ahci_device {
	int port_no;
	int type;
	volatile struct hba_port *port;

	struct hba_cmd_header cmd_header[32];
	char fis_data[256];
};

/************************** STOLEN **********************/

#define ATA_DEV_BUSY 0x80
#define ATA_DEV_DRQ 0x08

// Find a free command list slot
int find_cmdslot(volatile struct hba_port *port)
{
	// If not set in SACT and CI, the slot is free
	uint32_t slots = (port->sact | port->ci);
	for (int i=0; i<32; i++)
	{
		if ((slots&1) == 0)
			return i;
		slots >>= 1;
	}
	printk("Cannot find free command list entry");
	return -1;
}

static struct hba_cmd_table *create_cmd_table(struct iovec *vec, size_t nvec)
{
	size_t sz = sizeof(struct hba_cmd_table) + (sizeof(struct hba_prdt_entry) * nvec);
	struct hba_cmd_table *out = kzalloc(sz, VM_NORMAL);

	for (size_t i = 0; i < nvec; i++) {
		phys_addr_t vec_phys = vm_virt_to_phys(vec->io_buf);
		out->prdt_entry[i].dba = vec_phys & 0xFFFFFFFF;
		out->prdt_entry[i].dbau = (vec_phys >> 32) & 0xFFFFFFFF;
		out->prdt_entry[i].dbc = vec->io_len - 1;
		out->prdt_entry[i].i = (i == nvec - 1) ? 1 : 0;
	}

	return out;
}

static void free_cmd_table(struct hba_cmd_table *table)
{
	kfree(table);
}

static kern_status_t send_command(struct ahci_device *dev, unsigned int cmd, struct iovec *vec, size_t nvec)
{
	if (nvec == 0) {
		return KERN_OK;
	}

	volatile struct hba_port *port = dev->port;
	port->is = (uint32_t) -1;
	int spin = 0;
	int slot = find_cmdslot(port);
	if (slot == -1) {
		return KERN_BUSY;
	}

	struct hba_cmd_header *cmdheader = &dev->cmd_header[slot];
	cmdheader->cfl = sizeof(struct fis_reg_h2d) / sizeof(uint32_t);
	cmdheader->w = 0;
	cmdheader->prdtl = nvec;

	struct hba_cmd_table *cmdtbl = create_cmd_table(vec, nvec);
	phys_addr_t cmdtbl_phys = vm_virt_to_phys(cmdtbl);
	cmdheader->ctba = cmdtbl_phys & 0xFFFFFFFF;
	cmdheader->ctbau = (cmdtbl_phys >> 32) & 0xFFFFFFFF;

	struct fis_reg_h2d *cmdfis = (struct fis_reg_h2d*)(&cmdtbl->cfis);

	cmdfis->fis_type = FIS_TYPE_REG_H2D;
	cmdfis->c = 1;
	cmdfis->command = cmd;

	while ((port->tfd & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && spin < 1000000) {
		spin++;
	}

	if (spin == 1000000) {
		printk("Port is hung");
		free_cmd_table(cmdtbl);
		return KERN_DEVICE_STUCK;
	}

	port->ci = 1 << slot;

	// Wait for completion
	while (1)
	{
		if ((port->ci & (1 << slot)) == 0) {
			break;
		}

		if (port->is & HBA_PxIS_TFES) {
			printk("Read disk error 1");
			free_cmd_table(cmdtbl);
			return KERN_IO_ERROR;
		}
	}

	free_cmd_table(cmdtbl);

	if (port->is & HBA_PxIS_TFES) {
		printk("Read disk error 2");
		return KERN_IO_ERROR;
	}

	return KERN_OK;
}

static kern_status_t identify_device(struct ahci_device *dev, struct identify_device_data *out)
{
	/*
	port->is = (uint32_t) -1;		// Clear pending interrupt bits
	int spin = 0; // Spin lock timeout counter
	int slot = find_cmdslot(port);
	if (slot == -1)
		return false;

	struct hba_cmd_header *cmdheader = vm_phys_to_virt(port->clb);
	cmdheader += slot;
	cmdheader->cfl = sizeof(struct fis_reg_h2d)/sizeof(uint32_t);	// Command FIS size
	cmdheader->w = 0;		// Read from device
	cmdheader->prdtl = 1;

	struct hba_cmd_table *cmdtbl = vm_phys_to_virt(cmdheader->ctba);
	memset(cmdtbl, 0, sizeof(struct hba_cmd_table) +
 		(cmdheader->prdtl-1)*sizeof(struct hba_prdt_entry));

	cmdtbl->prdt_entry[0].dba = (uint32_t)vm_virt_to_phys(out);
	cmdtbl->prdt_entry[0].dbc = sizeof *out - 1;
	cmdtbl->prdt_entry[0].i = 1;

	// Setup command
	struct fis_reg_h2d *cmdfis = (struct fis_reg_h2d*)(&cmdtbl->cfis);

	cmdfis->fis_type = FIS_TYPE_REG_H2D;
	cmdfis->c = 1;	// Command
	cmdfis->command = ATA_CMD_IDENTIFY_DEVICE;

	// The below loop waits until the port is no longer busy before issuing a new command
	while ((port->tfd & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && spin < 1000000)
	{
		spin++;
	}
	if (spin == 1000000)
	{
		printk("Port is hung");
		return -1;
	}

	port->ci = 1<<slot;	// Issue command

	// Wait for completion
	while (1)
	{
		// In some longer duration reads, it may be helpful to spin on the DPS bit
		// in the PxIS port field as well (1 << 5)
		if ((port->ci & (1<<slot)) == 0)
			break;
		if (port->is & HBA_PxIS_TFES)	// Task file error
		{
			printk("Read disk error 1");
			return -1;
		}
	}

	// Check again
	if (port->is & HBA_PxIS_TFES)
	{
		printk("Read disk error 2");
		return -1;
	}

	return 0;
	*/
	struct iovec vec = { .io_buf = out, .io_len = sizeof *out };
	kern_status_t status = send_command(dev, ATA_CMD_IDENTIFY_DEVICE, &vec, 1);
	if (status != KERN_OK) {
		return status;
	}

	out->model_number[39] = 0;
	for (int i = 0; i < 40; i += 2) {
		unsigned char c = out->model_number[i];
		out->model_number[i] = out->model_number[i + 1];
		out->model_number[i + 1] = c;
	}

	for (int i = sizeof out->model_number - 1; i >= 0; i--) {
		if (isspace(out->model_number[i]) || out->model_number[i] == 0) {
			out->model_number[i] = 0;
		} else {
			break;
		}
	}

	return KERN_OK;
}

// Start command engine
void start_cmd(volatile struct hba_port *port)
{
	// Wait until CR (bit15) is cleared
	while (port->cmd & HBA_PxCMD_CR)
		;

	// Set FRE (bit4) and ST (bit0)
	port->cmd |= HBA_PxCMD_FRE;
	port->cmd |= HBA_PxCMD_ST;
}

// Stop command engine
void stop_cmd(volatile struct hba_port *port)
{
	// Clear ST (bit0)
	port->cmd &= ~HBA_PxCMD_ST;

	// Clear FRE (bit4)
	port->cmd &= ~HBA_PxCMD_FRE;

	// Wait until FR (bit14), CR (bit15) are cleared
	while(1)
	{
		if (port->cmd & HBA_PxCMD_FR)
			continue;
		if (port->cmd & HBA_PxCMD_CR)
			continue;
		break;
	}

}

void port_rebase(struct ahci_device *dev)
{
	volatile struct hba_port *port = dev->port;

	stop_cmd(port);	// Stop command engine

	struct hba_cmd_header *cmdheader = dev->cmd_header;

	// Command list offset: 1K*portno
	// Command list entry size = 32
	// Command list entry maxim count = 32
	// Command list maxim size = 32*32 = 1K per port
	phys_addr_t cmdheader_phys = vm_virt_to_phys(cmdheader);
	port->clb = cmdheader_phys & 0xFFFFFFFF;
	port->clbu = (cmdheader_phys >> 32) & 0xFFFFFFFF;
	memset(cmdheader, 0, sizeof *cmdheader);

	// FIS offset: 32K+256*portno
	// FIS entry size = 256 bytes per port
	phys_addr_t fis_phys = vm_virt_to_phys(dev->fis_data);
	port->fb = fis_phys & 0xFFFFFFFF;
	port->fbu = (fis_phys >> 32) & 0xFFFFFFFF;
	memset(dev->fis_data, 0, 256);

	/*
	// Command table offset: 40K + 8K*portno
	// Command table size = 256*32 = 8K per port
	for (int i=0; i<32; i++)
	{
		cmdheader[i].prdtl = 8;	// 8 prdt entries per command table
					// 256 bytes per command table, 64+16+48+16*8
		// Command table offset: 40K + 8K*portno + cmdheader_index*256
		cmdheader[i].ctba = phys + offset;
		cmdheader[i].ctbau = 0;
		memset(virt + offset, 0, 256);

		offset += 256;
	}
	*/

	start_cmd(port);	// Start command engine
}

/*
static int read(struct hba_port *port, uint32_t startl, uint32_t starth, uint32_t count, uint16_t *buf)
{
	port->is = (uint32_t) -1;		// Clear pending interrupt bits
	int spin = 0; // Spin lock timeout counter
	int slot = find_cmdslot(port);
	if (slot == -1)
		return false;

	struct hba_cmd_header *cmdheader = (struct hba_cmd_header*)(uintptr_t)port->clb;
	cmdheader += slot;
	cmdheader->cfl = sizeof(struct fis_reg_h2d)/sizeof(uint32_t);	// Command FIS size
	cmdheader->w = 0;		// Read from device
	cmdheader->prdtl = (uint16_t)((count-1)>>4) + 1;	// PRDT entries count

	struct hba_cmd_table *cmdtbl = (struct hba_cmd_table*)((uintptr_t)cmdheader->ctba);
	memset(cmdtbl, 0, sizeof(struct hba_cmd_table) +
 		(cmdheader->prdtl-1)*sizeof(struct hba_prdt_entry));

	// 8K bytes (16 sectors) per PRDT
	int i;
	for (i=0; i<cmdheader->prdtl-1; i++)
	{
		cmdtbl->prdt_entry[i].dba = (uint32_t)(uintptr_t)buf;
		cmdtbl->prdt_entry[i].dbc = 8*1024-1;	// 8K bytes (this value should always be set to 1 less than the actual value)
		cmdtbl->prdt_entry[i].i = 1;
		buf += 4*1024;	// 4K words
		count -= 16;	// 16 sectors
	}
	// Last entry
	cmdtbl->prdt_entry[i].dba = (uint32_t)(uintptr_t)buf;
	cmdtbl->prdt_entry[i].dbc = (count<<9)-1;	// 512 bytes per sector
	cmdtbl->prdt_entry[i].i = 1;

	// Setup command
	struct fis_reg_h2d *cmdfis = (struct fis_reg_h2d*)(&cmdtbl->cfis);

	cmdfis->fis_type = FIS_TYPE_REG_H2D;
	cmdfis->c = 1;	// Command
	cmdfis->command = ATA_CMD_READ_DMA_EX;

	cmdfis->lba0 = (uint8_t)startl;
	cmdfis->lba1 = (uint8_t)(startl>>8);
	cmdfis->lba2 = (uint8_t)(startl>>16);
	cmdfis->device = 1<<6;	// LBA mode

	cmdfis->lba3 = (uint8_t)(startl>>24);
	cmdfis->lba4 = (uint8_t)starth;
	cmdfis->lba5 = (uint8_t)(starth>>8);

	cmdfis->count_l = count & 0xFF;
	cmdfis->count_h = (count >> 8) & 0xFF;

	// The below loop waits until the port is no longer busy before issuing a new command
	while ((port->tfd & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && spin < 1000000)
	{
		spin++;
	}
	if (spin == 1000000)
	{
		printk("Port is hung");
		return -1;
	}

	port->ci = 1<<slot;	// Issue command

	// Wait for completion
	while (1)
	{
		// In some longer duration reads, it may be helpful to spin on the DPS bit
		// in the PxIS port field as well (1 << 5)
		if ((port->ci & (1<<slot)) == 0)
			break;
		if (port->is & HBA_PxIS_TFES)	// Task file error
		{
			printk("Read disk error 1");
			return -1;
		}
	}

	// Check again
	if (port->is & HBA_PxIS_TFES)
	{
		printk("Read disk error 2");
		return -1;
	}

	return 0;
}
*/

static int check_type(volatile struct hba_port *port)
{
	uint32_t ssts = port->ssts;

	uint8_t ipm = (ssts >> 8) & 0x0F;
	uint8_t det = ssts & 0x0F;

	if (det != HBA_PORT_DET_PRESENT) {
		return AHCI_DEV_NULL;
	}

	if (ipm != HBA_PORT_IPM_ACTIVE) {
		return AHCI_DEV_NULL;
	}

	switch (port->sig) {
	case SATA_SIG_ATAPI:
		return AHCI_DEV_SATAPI;
	case SATA_SIG_SEMB:
		return AHCI_DEV_SEMB;
	case SATA_SIG_PM:
		return AHCI_DEV_PM;
	default:
		return AHCI_DEV_SATA;
	}
}

void probe_ports(struct driver *driver, struct device *controller, volatile struct hba_config *abar)
{
	uint32_t pi = abar->pi;
	for (int i = 0; i < 32; i++) {
		if (!(pi & 1)) {
			pi >>= 1;
			continue;
		}

		int dt = check_type(&abar->ports[i]);

		switch (dt) {
		case AHCI_DEV_SATA:
		case AHCI_DEV_SATAPI:
			break;
		default:
			pi >>= 1;
			continue;
		}

		struct block_device *bdev = block_device_create();
		if (!bdev) {
			continue;
		}

		struct device *bdev_base = block_device_base(bdev);

		struct ahci_device *ahci_dev = kmalloc(sizeof *ahci_dev, VM_NORMAL);
		if (!ahci_dev) {
			device_deref(bdev_base);
			continue;
		}

		ahci_dev->port = &abar->ports[i];
		ahci_dev->port_no = i;
		ahci_dev->type = dt;
		bdev_base->dev_priv = ahci_dev;
		port_rebase(ahci_dev);

		snprintf(bdev_base->dev_name, sizeof bdev_base->dev_name, "ahci%d", i);

		if (dt == AHCI_DEV_SATA) {
			struct identify_device_data identity_data = {0};
			kern_status_t status = identify_device(ahci_dev, &identity_data);
			if (status != KERN_OK) {
				kfree(ahci_dev);
				device_deref(bdev_base);
				continue;
			}

			/*
			printk("ahci: found device '%s'", identity_data.model_number);
			printk("ahci: version: %u.%u", identity_data.major_revision, identity_data.major_revision);
			printk("ahci: C:%u, H:%u, S/T:%u", identity_data.num_cylinders, identity_data.num_heads, identity_data.num_sectors_per_track);
			printk("ahci: nr sectors=%u", identity_data.user_addressable_sectors);
			printk("ahci: offset: %u", offsetof(struct identify_device_data, current_sector_capacity));
			*/

			/* TODO read IDENTIFY DEVICE log to support non-512 sector sizes */
			bdev->sector_size = 512;
			bdev->capacity = identity_data.user_addressable_sectors;

			snprintf(bdev_base->dev_model_name, sizeof bdev_base->dev_model_name, "%s", identity_data.model_number);
		}

		device_register(bdev_base, driver, controller);

		pi >>= 1;
	}
}

/************************** STOLEN **********************/

static struct pci_driver *ahci_driver = NULL;
static volatile struct hba_config *hba_config = NULL;

static int irq_callback(void);

static struct irq_hook irq_hook = {
	.irq_callback = irq_callback,
};

static struct pci_device_id ahci_device_ids[] = {
	PCI_CLASS_ID(0x01, 0x06),
	PCI_DEVICE_ID_INVALID,
};


static kern_status_t ahci_scan(struct device *ahci_bus)
{
	printk("ahci: scanning connected SATA devices");
	uint32_t abar = pci_device_read_field(ahci_bus, PCI_REG_BAR5, 4);
	hba_config = vm_phys_to_virt(abar);
	probe_ports(pci_driver_base(ahci_driver), ahci_bus, hba_config);

	uint8_t irq = pci_device_read_field(ahci_bus, PCI_REG_INTERRUPT_LINE, 1);
	printk("ahci: IRQ line: %u", irq);
	return KERN_OK;
}

static struct bus_device_ops ahci_bus_ops = {
	.scan = ahci_scan,
};

static int irq_callback(void)
{
	printk("ahci: received IRQ");
	return 0;
}

static kern_status_t ahci_probe(struct pci_driver *driver, struct device *dev)
{
	printk("ahci: found AHCI controller");
	struct bus_device *ahci_bus = bus_device_from_generic(dev);
	ahci_bus->b_ops = &ahci_bus_ops;
	snprintf(dev->dev_name, sizeof dev->dev_name, "ahci");

	uint8_t irq_vec = pci_device_read_field(dev, PCI_REG_INTERRUPT_LINE, 1);
	hook_irq(IRQ0 + irq_vec, &irq_hook);

	return device_register(dev, pci_driver_base(ahci_driver), NULL);
}

static kern_status_t online(struct kext *self)
{
	printk("ahci: registering AHCI driver");
	ahci_driver = pci_driver_create(self, "ahci", ahci_device_ids);
	if (!ahci_driver) {
		return KERN_NO_MEMORY;
	}

	ahci_driver->probe = ahci_probe;

	pci_driver_register(ahci_driver);
	return KERN_OK;
}

DEFINE_KEXT("net.doorstuck.socks.ahci",
	    online, NULL,
	    PCI_SUBSYSTEM_KEXT_ID);