#include <stdint.h>
#include <stddef.h>
#include <vga.h>
#include <ahci.h>
#include <memory.h>
#include <mem.h>

enum fis_type {
	FIS_TYPE_REG_H2D	= 0x27, // Register FIS - host to device
	FIS_TYPE_REG_D2H	= 0x34, // Register FIS - device to host
	FIS_TYPE_DMA_ACT	= 0x39, // DMA activate FIS - device to host
	FIS_TYPE_DMA_SETUP	= 0x41, // DMA setup FIS - bidirectional
	FIS_TYPE_DATA		= 0x46, // Data FIS - bidirectional
	FIS_TYPE_BIST		= 0x58, // BIST activate FIS - bidirectional
	FIS_TYPE_PIO_SETUP	= 0x5f, // PIO setup FIS - device to host
	FIS_TYPE_DEV_BITS	= 0xa1 // Set device bits FIS - device to host
};

struct fis_reg_h2d {
	// DWORD 0
	uint8_t  fis_type; // FIS_TYPE_REG_H2D

	uint8_t  port_mul : 4; // Port multiplier
	uint8_t  rsv0 : 3;     // Reserved
	uint8_t  cmd_ctrl : 1; // 1: Command, 0: Control

	uint8_t  command;     // Command register
	uint8_t  feature_low; // Feature register, 7:0

	// DWORD 1
	uint8_t  lba0;       // LBA low register, 7:0
	uint8_t  lba1;       // LBA mid register, 15:8
	uint8_t  lba2;       // LBA high register, 23:16
	uint8_t  device_reg; // Device register

	// DWORD 2
	uint8_t  lba3;          // LBA register, 31:24
	uint8_t  lba4;          // LBA register, 39:32
	uint8_t  lba5;          // LBA register, 47:40
	uint8_t  feature_high;  // Feature register, 15:8

	// DWORD 3
	uint8_t  count_low;  // Count register, 7:0
	uint8_t  count_high; // Count register, 15:8
	uint8_t  icc;        // Isochronous command completion
	uint8_t  control;    // Control register

	// DWORD 4
	uint8_t  rsv1[4];    // Reserved
};

struct fis_reg_d2h {
	// DWORD 0
	uint8_t  fis_type; // FIS_TYPE_REG_D2H

	uint8_t  port_mul : 4;  // Port multiplier
	uint8_t  rsv0 : 2;      // Reserved
	uint8_t  interrupt : 1; // Interrupt bit
	uint8_t  rsv1 : 1;      // Reserved

	uint8_t  status; // Status register
	uint8_t  error;  // Error register

	// DWORD 1
	uint8_t  lba0;   // LBA low register, 7:0
	uint8_t  lba1;   // LBA mid register, 15:8
	uint8_t  lba2;   // LBA high register, 23:16
	uint8_t  device; // Device register

	// DWORD 2
	uint8_t  lba3;   // LBA register, 31:24
	uint8_t  lba4;   // LBA register, 39:32
	uint8_t  lba5;   // LBA register, 47:40
	uint8_t  rsv2;   // Reserved

	// DWORD 3
	uint8_t  count_low;  // Count register, 7:0
	uint8_t  count_high; // Count register, 15:8
	uint8_t  rsv3[2];    // Reserved

	// DWORD 4
	uint8_t  rsv4[4];     // Reserved
};

struct fis_data {
	// DWORD 0
	uint8_t fis_type;
	uint8_t port_mul : 4;

	uint8_t rsv0 : 4;
	uint8_t rsv1[2];

	// DWORD 1 ~ N
	uint32_t data[];
};

struct fis_pio_setup {
	uint8_t fis_type;
	uint8_t port_mul : 4;
	uint8_t rsv0 : 1;

	uint8_t direction : 1;
	uint8_t interrupt : 1;

	uint8_t rsv1 : 1;

	uint8_t status;
	uint8_t error;

	// DWORD 1
	uint8_t lba0;		// LBA low register, 7:0
	uint8_t lba1;		// LBA mid register, 15:8
	uint8_t lba2;		// LBA high register, 23:16
	uint8_t device;		// Device register

	// DWORD 2
	uint8_t lba3;		// LBA register, 31:24
	uint8_t lba4;		// LBA register, 39:32
	uint8_t lba5;		// LBA register, 47:40
	uint8_t rsv2;		// Reserved

	// DWORD 3
	uint8_t countl;		// Count register, 7:0
	uint8_t counth;		// Count register, 15:8
	uint8_t rsv3;		// Reserved
	uint8_t e_status;	// New value of status register

	// DWORD 4
	uint16_t tc;		// Transfer count
	uint8_t rsv4[2];	// Reserved
};

struct fis_dma_setup {
	// DWORD 0
	uint8_t  fis_type; // FIS_TYPE_DMA_SETUP

	uint8_t  port_mul : 4;      // Port multiplier
	uint8_t  rsv0 : 1;          // Reserved
	uint8_t  direction : 1;     // Data transfer direction, 1 - device to host
	uint8_t  interrupt : 1;     // Interrupt bit
	uint8_t  auto_activate : 1; // Auto-activate. Specifies if DMA Activate FIS is needed

	uint8_t  rsved[2]; // Reserved

	//DWORD 1&2

	uint64_t DMAbufferID; // DMA Buffer Identifier. Used to Identify DMA buffer in host memory.
							 // SATA Spec says host specific and not in Spec. Trying AHCI spec might work.

	//DWORD 3
	uint32_t rsv1; //More reserved

	//DWORD 4
	uint32_t DMAbufOffset; //Byte offset into buffer. First 2 bits must be 0

	//DWORD 5
	uint32_t TransferCount; //Number of bytes to transfer. Bit 0 must be 0

	//DWORD 6
	uint32_t rsv2; //Reserved
};

struct hba_cmd_header {
	// DW0
	uint8_t  cfl : 5;  // Command FIS length in DWORDS, 2 ~ 16
	uint8_t  a : 1;    // ATAPI
	uint8_t  w : 1;    // Write, 1: H2D, 0: D2H
	uint8_t  p : 1;    // Prefetchable
 
	uint8_t  r : 1;    // Reset
	uint8_t  b : 1;    // BIST
	uint8_t  c : 1;    // Clear busy upon R_OK
	uint8_t  rsv0 : 1; // Reserved
	uint8_t  pmp : 4;  // Port multiplier port
 
	uint16_t prdtl; // Physical region descriptor table length in entries
 
	// DW1
	volatile
	uint32_t prdbc; // Physical region descriptor byte count transferred
 
	// DW2, 3
	uint32_t ctba;  // Command table descriptor base address
	uint32_t ctbau; // Command table descriptor base address upper 32 bits
 
	// DW4 - 7
	uint32_t rsv1[4]; // Reserved
};

struct hba_prdt_entry {
	uint32_t data_base_addr;
	uint32_t data_base_addr_upper;
	uint32_t reserved_0;

	uint32_t data_byte_count : 22;
	uint32_t reserved_1 : 9;
	uint32_t interrupt : 1;
};

struct hba_cmd_table {
	uint8_t command_fis[64];
	uint8_t atapi_command[16];
	uint8_t reserved[48];
	struct hba_prdt_entry entries[];
};

enum ahci_dev {
	AHCI_DEV_NULL = 0,
	AHCI_DEV_SATA = 1,
	AHCI_DEV_SEMB = 2,
	AHCI_DEV_PM = 3,
	AHCI_DEV_SATAPI = 4
};

// TODO: replace with enum when c23 underlying types are implemented
enum foo : size_t {
	SATA_SIG_ATA = 0x00000101,   // SATA drive
	SATA_SIG_ATAPI = 0xEB140101, // SATAPI drive
	SATA_SIG_SEMB = 0xC33C0101,  // Enclosure management bridge
	SATA_SIG_PM = 0x96690101     // Port multiplier
};

struct disk disks[5];

struct disk *get_disk(int id) {
	return &disks[id];
}

static enum ahci_dev check_type(struct ahci_port *port) {
	uint32_t ssts = port->ssts;
	if ((ssts & 0xf) != 0x3)
		return AHCI_DEV_NULL;
	if (((ssts >> 8) & 0xf) != 0x1)
		return AHCI_DEV_NULL;

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

void start_cmd(struct ahci_port *port) {
	//while (port->cmd & (1 << 15));
	while (port->cmd & 0x8000);
	//port->cmd |= (1 << 1) | (1 << 0);
	port->cmd |= 0x0010;
	port->cmd |= 0x0001;
}

void stop_cmd(struct ahci_port *port) {
	//port->cmd &= ~(1 << 0) | ~(1 << 1);
	port->cmd &= ~0x0001;
	port->cmd &= ~0x0010;

	//while (port->cmd & ((1 << 15) | (1 << 14)));
	while (1) {
		if (port->cmd & 0x4000)
			continue;
		if (port->cmd & 0x8000)
			continue;
		break;
	}
}

static void port_alloc(struct ahci_port *port) {
	stop_cmd(port);

	struct hba_cmd_header *cmd_header = kmalloc_a(1024, 1024);
	port->clb = (uintptr_t) cmd_header & 0xffffffff;
	port->clbu = (uintptr_t) cmd_header >> 32;
	memset(cmd_header, 0, 1024);

	void *fb = kmalloc_a(256, 256);
	port->fb = (uintptr_t) fb & 0xffffffff;
	port->fbu = (uintptr_t) fb >> 32;
	memset(fb, 0, 256);

	for (size_t i = 0; i < 32; i++) {
		cmd_header[i].prdtl = 8;
		void *ctba = kmalloc_a(256, 128);
		cmd_header[i].ctba = (uintptr_t) ctba & 0xffffffff;
		cmd_header[i].ctbau = (uintptr_t) ctba >> 32;
		memset(ctba, 0, 256);
	}

	start_cmd(port);
}

int32_t find_cmdslot(volatile struct ahci_port *port) {
	uint32_t slots = (port->sact | port->ci);
	for (size_t i = 0; i < 32; i++)
		if ((slots & (1 << i)) == 0)
			return i;
	return -1;
}

static void load_ptrd(uint16_t len, struct hba_prdt_entry entries[len],
		uint16_t *buffer, uint32_t count) {
	for (uint16_t i = 0; i < len - 1; i++) {
		entries[i] = (struct hba_prdt_entry) {
			.data_base_addr = (uintptr_t) buffer & 0xffffffff,
			.data_base_addr_upper = (uintptr_t) buffer >> 32,
			.data_byte_count = 8 * 1024 - 1,
			.interrupt = 1
		};
		buffer += 4 * 1024;
		count -= 16;
	}

	entries[len - 1] = (struct hba_prdt_entry) {
		.data_base_addr = (uintptr_t) buffer & 0xffffffff,
		.data_base_addr_upper = (uintptr_t) buffer >> 32,
		.data_byte_count = (count << 9) - 1,
		.interrupt = 1
	};
}

bool read(volatile struct ahci_port *port, uint64_t offset, uint32_t count, uint16_t buffer[256 * count]) {
	port->is = -1;
	size_t spin = 0;
	int32_t slot = find_cmdslot(port);
	if (slot == -1)
		return false;

	struct hba_cmd_header *cmd_header = (void *)(((uintptr_t)port->clbu << 32) | port->clb);
	cmd_header += slot;
	cmd_header->cfl = sizeof(struct fis_reg_h2d) / sizeof(uint32_t);
	cmd_header->w = 0;
	cmd_header->prdtl = (uint16_t)((count - 1) >> 4) + 1;

	struct hba_cmd_table *cmd_table = (void *)((uintptr_t)cmd_header->ctbau << 32 | cmd_header->ctba);
	*cmd_table = (struct hba_cmd_table) {0};
	load_ptrd(cmd_header->prdtl, cmd_table->entries, buffer, count);

	struct fis_reg_h2d *cmd_fis = (void *) &cmd_table->command_fis;
	*cmd_fis = (struct fis_reg_h2d) {
		.fis_type = FIS_TYPE_REG_H2D,
		.cmd_ctrl = 1,
		.command = 0x25,
		.lba0 = (uint8_t) offset & 0xff,
		.lba1 = (uint8_t) (offset >> 8) & 0xff,
		.lba2 = (uint8_t) (offset >> 16) & 0xff,
		.lba3 = (uint8_t) (offset >> 24) & 0xff,
		.lba4 = (uint8_t) (offset >> 32) & 0xff,
		.lba5 = (uint8_t) (offset >> 40) & 0xff,
		.device_reg = 1 << 6,
		.count_low = count & 0xff,
		.count_high = (count >> 8) & 0xff
	};

	while ((port->tfd & (0x80 | 0x08)) && spin < 10000000)
		spin++;

	if (spin == 10000000)
		return false;

	port->ci = 1 << slot;

#define ATA_DEV_BUSY 0x80
#define ATA_DEV_DRQ 0x08

	//while ((port->ci & (1 << slot)));
	while (1) {
		if ((port->ci & (1 << slot)) == 0)
			break;
		if (port->is & (1 << 30)) {
			vga_puts("disk read error\n");
			return false;
		}
	}

	return true;
}

static void identify(struct ahci_port *port) {
	port->is = -1;
	size_t spin = 0;
	int32_t slot = find_cmdslot(port);
	if (slot == -1)
		return;
	struct hba_cmd_header *cmd_header = &((struct hba_cmd_header*)(((uintptr_t)port->clbu << 32) | port->clb))[slot];
	cmd_header->cfl = sizeof(struct fis_reg_h2d) / sizeof(uint32_t);
	cmd_header->w = 0;
	cmd_header->prdtl = 1;

	uint16_t buf[256] = {0};
	struct hba_cmd_table *cmd_table = (void *)((uintptr_t)cmd_header->ctbau << 32 | cmd_header->ctba);
	*cmd_table = (struct hba_cmd_table) {0};

	load_ptrd(1, cmd_table->entries, buf, 1);

	struct fis_reg_h2d *cmd_fis = (void *) &cmd_table->command_fis;
	*cmd_fis = (struct fis_reg_h2d) {
		.fis_type = FIS_TYPE_REG_H2D,
		.cmd_ctrl = 1,
		.command = 0xec
	};

	while ((port->tfd & (0x80 | 0x08)) && spin < 10000000)
		spin++;

	if (spin == 10000000)
		return;

	port->ci = 1 << slot;

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

	uint64_t blocks = (uint64_t)buf[103] << 48 | (uint64_t)buf[102] << 32 | (uint64_t)buf[101] << 16 | buf[100];
	disks[0] = (struct disk) {
		.id = 0,
		.port = port,
		.block_count = blocks
	};
}

void probe_port(struct hba_mem *abar) {
	for (size_t i = 0; i < 32; i++) {
		if (!(abar->pi & (1 << i)))
			continue;

		memory_map((uintptr_t)&abar->ports[i], sizeof(struct ahci_port));

		const char *drive;
		switch (check_type(&abar->ports[i])) {
			case AHCI_DEV_SATA:
				drive = "SATA";
				port_alloc(&abar->ports[i]);
				identify(&abar->ports[i]);
				break;
			case AHCI_DEV_SATAPI:
				drive = "SATAPI";
				break;
			case AHCI_DEV_SEMB:
				drive = "SEMB";
				break;
			case AHCI_DEV_PM:
				drive = "PM";
				break;
			case AHCI_DEV_NULL:
				drive = "No";
				break;
		}
		vga_puts(drive);
		vga_puts(" drive found at port ");
		vga_putn(i);
		vga_putc('\n');
	}
}