/**
 * Contains functions for the AHCI driver
 */

#include <kernel/kernel.h>
#include <drv/ahci.h>
#include <kernel/slab.h>
#include <kernel/paging.h>
#include <kernel/heap.h>
#include <lib/string.h>
#include <kernel/idt.h>
#include <kernel/block.h>

/* Pointer to the current AHCI controllers HBA registers */
AHCI_HBA_MEMORY *hba_mem;

/**
 * Rebase AHCI memory to the memory address given by AHCI_START
 *
 * params:
 *		mem - AHCI controller memory of the controller to rebase.
 */
void ahci_rebase(AHCI_HBA_MEMORY *mem);

/**
 * Start the command engine on a port.
 *
 * params:
 *		port - Port to start the engine on.
 */
void ahci_start_cmd(AHCI_HBA_PORT *port);


/**
 * Stop the command engine on a port.
 *
 * params:
 *		port - Port to stop the engine on.
 */
void ahci_stop_cmd(AHCI_HBA_PORT *port);

/**
 * Probe AHCI ports for disks and initialize each disk.
 *
 * params:
 *		mem - AHCI HBA memory registers for the controller.
 */
void ahci_find_disks(AHCI_HBA_MEMORY *mem);

/**
 * Find an unused command slot in an AHCI HBA port.
 *
 * params:
 * port - Port to search for a command slot on.
 */
uint8_t ahci_find_command_slot(AHCI_HBA_PORT *port);

/**
 * Read/write sectors from the initialized AHCI device.
 *
 * params:
 *		port - Port on which a disk is attached.
 *		lba_low - Lower 32 bits of LBA.
 *		lba_high - Higher 32 bits of LBA.
 *		num_sectors - Number of sectors to read.
 *		buf - Where to write data to.
 *		write - ATA opcode to send.
 */
void ahci_send_command(AHCI_HBA_PORT *port, uint32_t lba_low, uint32_t lba_high, uint32_t num_sectors, uint32_t buf, uint8_t command);

/**
 * Read sectors from a disk controlled by AHCI.
 *
 * params:
 *		dev - Block device describing the AHCI port.
 *		req - The request to be fulfilled.
 */
void ahci_post_request(BLOCK_DEVICE *dev, BLOCK_DEVICE_REQUEST *req);

/**
 * Issue a command FIS to a slot.
 *
 * params:
 *		port - Port to issue the command to.
 *		slot - Slot to issue the command to.
 */
void ahci_issue_command(AHCI_HBA_PORT *port, uint8_t slot);

/**
 * Interrupt handler for the AHCI driver.
 *
 * params:
 *		registers - The register state at interrupt time.
 */
void ahci_isr(ISR_REG *registers);

/**
 * Perform a software reset on a device connect to a port.
 *
 * params:
 *		port - The port to send a software reset to.
 */
void ahci_reset_port(AHCI_HBA_PORT *port);

/**
 * Initialize AHCI structures.
 */
void ahci_init()
{
	hba_mem = NULL;
}

void ahci_device_init(PCI_DEVICE *gdev)
{	
	/* Sorry, but only one AHCI controller only right now. If HBA_MEM is initialized, get out of here. */
	if (hba_mem)
	{
		return;
	}

	/* Enable interrupts, DMA, and memory space access from the PCI command register */
	uint32_t cmd = pci_config_readl(gdev->common_header.bus, gdev->common_header.device, gdev->common_header.function, PCI_COMMAND);
	cmd &= ~(0x1 << 9); /* Ensure interrupt disable is not set */
	cmd |= (0x3 << 1);  /* Enable bus mastering and memory space access */
	pci_config_writel(gdev->common_header.bus, gdev->common_header.device, gdev->common_header.function, PCI_COMMAND, cmd);
	
	/* Memory map the AHCI BAR as uncacheable */
	for (uint32_t i = 0; i < gdev->bar[5].limit; i += 0x1000)
	{
		map_page(gdev->bar[5].base + i, gdev->bar[5].base + i, get_kernel_pd(), PTE_RW | PTE_PCD);
	}
	
	/* Memory map AHCI memory as uncacheable */
	for (uint32_t i = 0; i < AHCI_SIZE; i += 0x1000)
	{
		map_page(AHCI_START + i, AHCI_START + i, get_kernel_pd(), PTE_RW | PTE_PCD);
	}
	
	/* Get the HBA memory registers */
	hba_mem = (AHCI_HBA_MEMORY *) (gdev->bar[5].base);
	
	/* Check if this device is capable of a BIOS/OS handoff. If it is, make sure we have ownership. If we don't, take ownership. */
	if (hba_mem->cap2 & CAP2_BOH)
	{
		/* If BOHC_OOS is 0 and BOHC_BOS is 0, we have not obtained ownership */
		if (!(hba_mem->bohc & BOHC_OOS) && (hba_mem->bohc & BOHC_BOS))
		{
			/* Per the spec, set BOHC_OOS to 1, and spin on BOHC_BOS until it is 0 */
			hba_mem->bohc |= BOHC_OOS;
			
			uint32_t spin = 0;
			while (hba_mem->bohc & BOHC_BOS)
			{
				if (spin > 1000000)
				{
					kprintf("BIOS/OS handoff was not completed. There may be problems now...\n");
					break;
				}
				else
				{
					spin++;
				}
			}
		}
	}
	
	/* Reset the controller */
	hba_mem->ghc |= GHC_HR;
	
	/* Spin while the controller resets */
	while (hba_mem->ghc & GHC_HR) {}
	
	/* Plug in the interrupt handler */
	register_interrupt_handler(gdev->common_header.interrupt_line, &ahci_isr);
	
	/* Enable AHCI mode and interrupts in the GHC register */
	hba_mem->ghc |= (GHC_AE | GHC_IE);

	/* Rebase our AHCI structures */
	ahci_rebase(hba_mem);
	
	/* Find and initialize disks */
	ahci_find_disks(hba_mem);
}

void ahci_start_cmd(AHCI_HBA_PORT *port)
{
	/* Spin until no commands are running. We can do this by checking bit 15 of the CMD register (CR) */
	while (port->cmd & PxCMD_CR) {}
	
	/* Stop the command engine by setting the FRE bit of the cmd register (bit 4) and ST (bit 0) */
	port->cmd |= PxCMD_FRE;
	port->cmd |= PxCMD_ST;
}

void ahci_stop_cmd(AHCI_HBA_PORT *port)
{
	/* Opposite of CMD_START */
	port->cmd &= ~PxCMD_ST;
	port->cmd &= ~PxCMD_FRE;
	
	/* Wait until FR and CR are cleared */
	while ((port->cmd & PxCMD_FR) || (port->cmd & PxCMD_CR)) {}
}

/* Yeah, it looks like I'm going to eat a big chunk of memory even if the AHCI controller doesn't support
 * all 32 ports. Oh well!
 */
void ahci_rebase(AHCI_HBA_MEMORY *mem)
{
	/* Enumerate over all ports */
	for (uint8_t port = 0; port < AHCI_NUM_PORTS; port++)
	{
		/* Check if the port is implemented */
		if (mem->pi & (0x1 << port))
		{
			/* It's implemented, we're good to go! */
			ahci_stop_cmd(&(mem->ports[port]));
			
			/* Command headers first. 32 headers * 32 bytes = 1KiB per port. 32KiB total. */
			mem->ports[port].clb = AHCI_START + (port << 10);
			
			/* FIS next. Size 256 * 32 ports = 8KiB total. */
			mem->ports[port].fb = AHCI_START + (AHCI_NUM_PORTS << 10) + (port << 8);
			
			/* Command table next. Size 256 * 32 command headers is 8KiB per port. */
			/* Each command table has eight PRDTs, each of size 256 */
			AHCI_COMMAND_HEADER *hdr = (AHCI_COMMAND_HEADER *) (mem->ports[port].clb);
			for (uint8_t i = 0; i < AHCI_COMMAND_LIST_SIZE; i++)
			{
				hdr[i].prdtl = AHCI_PRDT_SIZE; /* Set the PRDT length */
				hdr[i].ctba = AHCI_START + (40 << 10) + (port << 13) + (i << 8);
			}
			
			ahci_start_cmd(&(mem->ports[port]));
		}
	}
}

void ahci_find_disks(AHCI_HBA_MEMORY *mem)
{
	/* Iterate over all valid ports */
	for (uint8_t port = 0; port < AHCI_NUM_PORTS; port++)
	{
		if (mem->pi & (0x1 << port))
		{
			/* Check the IPM and DET of the ssts register. If the device and PHY is present, and the device is active, then it is valid */
			if ((mem->ports[port].ssts & PxSSTS_DET_DEV_PHY) && ((mem->ports[port].ssts >> 8) & PxSSTS_PM_ACTIVE))
			{
				switch (mem->ports[port].sig)
				{
					case SATA_SIG_ATAPI: /* ATAPI Disk */
					{
						// Initialize a disk
						kprintf("Drive found at port 1%u\n", port);
						break;
					}
					case SATA_SIG_SEMB:
					{
						// Initialize a disk
						kprintf("Drive found at port 2%u\n", port);
						break;
					}
					case SATA_SIG_PM:
					{
						// Initialize a disk
						kprintf("Drive found at port 3%u\n", port);
						break;
					}
					default: /* ATA disk */
					{
						/* Reset the port. See 10.4.1 of the AHCI spec. */
						ahci_reset_port(&(mem->ports[port]));
						
						/* Enable command list processing. PxCMD.ST set to 1 */
						mem->ports[port].cmd |= PxCMD_ST;
						
						/* Enable interrupts on this port */
						/* TODO: For whatever reason, setting this value to PxIE_DHRE, among other values, causes hanging. No clue why */
						/* Setting either PxIE_DHRE or PxIE_PSE will cause this to hang */
						mem->ports[port].ie = PxIE_SDBE;
						kprintf("%x\n", (uint32_t) &(mem->ports[port].ie));
						
						/* Set up a buffer to get disk info from an ATA_IDENTIFY_DEVICE command */
						uint16_t info[256];
						memset(info, 0, 512);
						
						/* Send IDENTIFY_DEVICE to get disk information */
						ahci_send_command(&(hba_mem->ports[port]), 0, 0, 1, (uint32_t) info, ATA_CMD_IDENTIFY_PIO);
						
						BLOCK_DEVICE dev = 
						{
							.name = "AHCI",
							.driver_type = AHCI,
							.port = port,
							.post_request = &ahci_post_request
						};
						
						/* Does this device support LBA48? */
						dev.lba48 = (info[63] >> 3) & 0x1;
						
						/* Set the number of blocks */
						for (uint8_t i = 100; i <= 103; i++)
						{
							dev.num_blocks[i] = info[i];
						}
						
						dev.block_size = 512;
						
						register_block_device(&dev);
					}
				}
			}
		}
	}
}

void ahci_reset_port(AHCI_HBA_PORT *port)
{
	/* To do comreset:
		0. Clear PxCMD.ST and wait for PxCMD.CR to clear.
		1. Write 0x1 to PxSCTL.DET
		2. Wait 1 millisecond
		3. Write 0x0 to PxSCTL.DET
		4. Spin on PxSSTS.DET being set to 0x3
		5. Write all 1s to PxSERR register
	*/
	uint32_t spin = 0;
	
	/* Clear PxCMD.ST and spin on PxCMD.CR */
	port->cmd &= ~(PxCMD_ST);
	
	/* If we spin too long, we should just proceed with the reset anyway */
	while ((port->cmd & PxCMD_CR) && (spin < 1000000))
	{
		spin++;
	}
	spin = 0;
	
	/* Set PxSCTL.Det to 0x1 */
	port->sctl &= ((~0x0) << 4);
	port->sctl |= 0x1;
	
	while (spin <= 1000000)
	{
		spin++;
	}
	
	/* This is a weird way of setting the first four bits of sctl to 0x0 */
	port->sctl &= ((~0x0) << 4);
	
	/* Set PxSCTL.Det back to 0x0 */
	while ((port->ssts & PxSSTS_DET) != 0x3) {}
	
	port->serr = ~0x0;
}

uint8_t ahci_find_command_slot(AHCI_HBA_PORT *port)
{
	/* A command slot is free if it is not set in PxSACT or PxCI */
	for (uint8_t i = 0; i < 32; i++)
	{
		if (!(port->sact & (0x1 << i)) && !(port->ci & (0x1 << i)))
		{
			return i;
		}
	}
	
	/* No command slot found */
	return 0xFF;
}

void ahci_issue_command(AHCI_HBA_PORT *port, uint8_t slot)
{
	/* Wait for the port to become clear */
	uint32_t spin = 0;
	while ((port->tfd & (PxTFD_STS_BSY | PxTFD_STS_DRQ)) && spin < 1000000)
	{
		spin++;
	}
	
	if (spin >= 1000000)
	{
		return;
	}
	
	port->ci = 1 << slot;
	
	AHCI_COMMAND_HEADER *hdr = (AHCI_COMMAND_HEADER *) (port->clb);
	AHCI_COMMAND_TABLE *tbl = (AHCI_COMMAND_TABLE *) hdr[slot].ctba;
	AHCI_FIS_H2D *fis = (AHCI_FIS_H2D *) (&tbl->cfis);
	
	/* Since identify is only available via PIO, we shouldn't spin on the port to wait for it to complete. */
	if (!(fis->command == ATA_CMD_IDENTIFY_PIO))
	{
		/* Wait for the command to complete */
		while ((port->ci) & (1 << slot))
		{
			/* Check for a task file error */
			if ((port->is) & PxIS_TFES)
			{
				/* Something bad happened */
				return;
			}
		}
		
		/* When we complete, check one more time for an error */
		if ((port->is) & PxIS_TFES)
		{
			/* Something bad happened */
			return;
		}
	}
}

void ahci_send_command(AHCI_HBA_PORT *port, uint32_t lba_low, uint32_t lba_high, uint32_t num_sectors, uint32_t buf, uint8_t command)
{
	/* Clear pending interrupt bits */
	port->is = ~0x0;
	
	/* Get an empty command slot */
	uint8_t slot = ahci_find_command_slot(port);
	
	AHCI_COMMAND_HEADER *hdr = (AHCI_COMMAND_HEADER *) (port->clb);
	/* CFL is the command FIS length *** IN TERMS OF DOUBLE WORDS ***. Divide by 4 to get the length we need. */
	hdr[slot].cfl = (sizeof(AHCI_FIS_H2D) >> 2);
	
	/* Set to 0 if reading, else 1. */
	hdr[slot].w = (command == ATA_CMD_WRITE_DMA_EX ? 1 : 0);
	
	/* PRDT length. One PRDT can do up to 4MiB, or 16 sectors */
	hdr[slot].prdtl = (uint16_t) (((num_sectors - 1) >> 4) + 1);
	if (hdr[slot].prdtl > 8)
	{
		/* Maximum PRDTLs is 8, so we can only read/write 32MiB at a time */
		return;
	}
	
	/* Get the command table */
	AHCI_COMMAND_TABLE *tbl = (AHCI_COMMAND_TABLE *) hdr[slot].ctba;
	memset(tbl, 0, sizeof(AHCI_COMMAND_TABLE));
	
	/* Set up our PRDTs. 4MiB per PRDT. */
	uint16_t i = 0;
	for (i; i < (hdr[slot].prdtl - 1); i++)
	{
		tbl->prdt_entry[i].dba = buf;			/* Tell the PRDT to write at buf */
		tbl->prdt_entry[i].dbc = (4 << 20) - 1;	/* 4MiB. Subtract one so that byte 0 is 1, per the AHCI spec */
		tbl->prdt_entry[i].i = 1;				/* Yes, please interrupt the CPU on completion. */
		
		buf += (4 << 20) >> 2;					/* Adjust buf to the right by 4MiB. Divide by two since the arithmetic will be in terms of DWORDs. */
		num_sectors -= 16;						/* We will have moved 16 sectors, so remove these from the sector count */
	}
	
	/* Last entry */
	tbl->prdt_entry[i].dba = buf;
	tbl->prdt_entry[i].dbc = (num_sectors << 9) - 1; /* Multiply number of sectors (remaining) by 512 for the final byte count */
	tbl->prdt_entry[i].i = 1;
	
	/* Setup the command */
	AHCI_FIS_H2D *fis = (AHCI_FIS_H2D *) (&tbl->cfis);
	
	fis->fis_type = HOST_TO_DEVICE;
	fis->c = 1;
	fis->command = command;
	
	fis->lba0 = (uint8_t) lba_low;
	fis->lba1 = (uint8_t) (lba_low >> 8);
	fis->lba2 = (uint8_t) (lba_low >> 16);
	fis->device = 1 << 6; /* Set LBA mode */
	
	fis->lba3 = (uint8_t) (lba_low >> 24);
	fis->lba4 = (uint8_t) lba_high;
	fis->lba5 = (uint8_t) (lba_high >> 8);
	
	fis->countl = num_sectors & 0xFF;
	fis->counth = (num_sectors >> 8) & 0xFF;

	ahci_issue_command(port, slot);
}

void ahci_post_request(BLOCK_DEVICE *dev, BLOCK_DEVICE_REQUEST *req)
{
	if (dev->driver_type == AHCI)
	{
		ahci_send_command(&(hba_mem->ports[dev->port]), req->lba_low, req->lba_high, req->num_sectors, req->buffer, (req->write ? ATA_CMD_WRITE_DMA_EX : ATA_CMD_READ_DMA_EX));
	}
}

void ahci_isr(ISR_REG *registers)
{
	kprintf("TODO: Write the AHCI ISR\n");
}