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

int find_cmdslot(volatile struct hba_port *port)
{
	uint32_t slots = (port->sact | port->ci);
	for (int i = 0; i < 32; i++)
	{
		if ((slots & 0x01u) == 0) {
			return i;
		}

		slots >>= 1;
	}

	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);
}

kern_status_t send_ata_command(struct ahci_device *dev, unsigned int cmd, struct iovec *vec, size_t nvec)
{
	return send_ata_command_ex(dev, NULL, cmd, vec, nvec);
}

kern_status_t send_ata_command_ex(struct ahci_device *dev, struct fis_reg_h2d *fis, 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->a = 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);

	if (fis) {
		memcpy(cmdfis, fis, sizeof *cmdfis);
	} else {
		memset(cmdfis, 0x00, sizeof *cmdfis);
	}

	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) {
		free_cmd_table(cmdtbl);
		return KERN_DEVICE_STUCK;
	}

	port->ci = 1 << slot;

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

		if (port->is & HBA_PxIS_TFES) {
			free_cmd_table(cmdtbl);
			printk("ahci: command failed");
			return KERN_IO_ERROR;
		}
	}
	free_cmd_table(cmdtbl);

	if (port->is & HBA_PxIS_TFES) {
		printk("ahci: command failed");
		return KERN_IO_ERROR;
	}

	return KERN_OK;
}

kern_status_t send_atapi_command(struct ahci_device *dev, struct scsi_command *cmd, struct iovec *vec, size_t nvec)
{
	return send_atapi_command_ex(dev, NULL, cmd, vec, nvec);
}

kern_status_t send_atapi_command_ex(struct ahci_device *dev, struct fis_reg_h2d *fis, struct scsi_command *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];
	memset(cmdheader, 0x0, sizeof *cmdheader);
	cmdheader->cfl = sizeof(struct fis_reg_h2d) / sizeof(uint32_t);
	cmdheader->w = 0;
	cmdheader->a = 1;
	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;
	memcpy(&cmdtbl->acmd, cmd, sizeof *cmd);

	struct fis_reg_h2d *cmdfis = (struct fis_reg_h2d *)(&cmdtbl->cfis);
	if (fis) {
		memcpy(cmdfis, fis, sizeof *cmdfis);
	} else {
		memset(cmdfis, 0x0, sizeof *cmdfis);
	}

	cmdfis->fis_type = FIS_TYPE_REG_H2D;
	cmdfis->c = 1;
	cmdfis->rsv0 = 0;
	cmdfis->feature_l = 1;
	cmdfis->command = ATA_CMD_PACKET;

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

	if (spin == 1000000) {
		free_cmd_table(cmdtbl);
		return KERN_DEVICE_STUCK;
	}

	port->ci = 1 << slot;

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

		if (port->is & HBA_PxIS_TFES) {
			free_cmd_table(cmdtbl);
			return KERN_IO_ERROR;
		}
	}

	volatile int _i = 0;
	while ((port->is & 0x01u) != 0x1) {
		_i++;
	}

	free_cmd_table(cmdtbl);

	if (port->is & HBA_PxIS_TFES) {
		return KERN_IO_ERROR;
	}

	return KERN_OK;
}

kern_status_t identify_ata_device(struct ahci_device *dev, struct identify_device_data *out)
{
	struct iovec vec = { .io_buf = out, .io_len = sizeof *out };
	kern_status_t status = send_ata_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;
}

kern_status_t identify_atapi_device(struct ahci_device *dev, struct identify_device_data *out)
{
	struct iovec vec = { .io_buf = out, .io_len = sizeof *out };
	kern_status_t status = send_ata_command(dev, ATA_CMD_IDENTIFY_PACKET_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)
{
	volatile int _i = 0;
	while (port->cmd & HBA_PxCMD_CR) {
		_i++;
	}

	port->cmd |= HBA_PxCMD_FRE;
	port->cmd |= HBA_PxCMD_ST;
}

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

	while(1) {
		if (port->cmd & HBA_PxCMD_FR) {
			continue;
		}

		if (port->cmd & HBA_PxCMD_CR) {
			continue;
		}

		break;
	}
}

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

	stop_cmd(port);

	struct hba_cmd_header *cmdheader = dev->cmd_header;

	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);

	phys_addr_t fis_phys = vm_virt_to_phys(dev->rfis_data);
	port->fb = fis_phys & 0xFFFFFFFF;
	port->fbu = (fis_phys >> 32) & 0xFFFFFFFF;
	memset(dev->rfis_data, 0, 256);

	start_cmd(port);
}

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_ahci_ports(struct driver *driver, struct device *controller, volatile struct hba_config *abar,
                      void(*callback)(int, struct device *, volatile struct hba_port *, int))
{
	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;
		}

		callback(i, controller, &abar->ports[i], dt);

		pi >>= 1;
	}
}