#include <sys.h>

#include "ata.h"
#include "pci.h"
#include "../panic.h"

#define REG_DATA 0
#define REG_ERR 1
#define REG_FEATURES 1
#define REG_SECTOR_CNT 2
#define REG_SECTOR_NUM 3
#define REG_LBA_LO 3
#define REG_CYLINDER_LO 4
#define REG_LBA_MID 4
#define REG_CYLINDER_HI 5
#define REG_LBA_HI 5
#define REG_DRIVE_HEAD 6
#define REG_STATUS 7
#define REG_COMMAND 7

#define STAT_ERR 0x01
#define STAT_IDX 0x02
#define STAT_CORR 0x04
#define STAT_DRQ 0x08
#define STAT_SRV 0x10
#define STAT_DF 0x20
#define STAT_RDY 0x40
#define STAT_BSY 0x80

static bool ata_init_impl(uint16_t base, uint16_t ctrl, struct AtaDevice *dev) {
	uint8_t status = inb(base + REG_STATUS);
	outb(base + REG_DRIVE_HEAD, 0xA0);
	outb(base + REG_SECTOR_CNT, 0x00);
	outb(base + REG_LBA_LO, 0x00);
	outb(base + REG_LBA_MID, 0x00);
	outb(base + REG_LBA_HI, 0x00);
	outb(base + REG_COMMAND, 0xEC);
	status = STAT_BSY;
	while(status & STAT_BSY) {
		status = inb(base + REG_STATUS);
	}
	if(inb(base + REG_LBA_MID) != 0 || inb(base + REG_LBA_HI) != 0) {
		panic("Device is not an ATA device");
	}
	while(!(status & (STAT_DRQ | STAT_ERR))) {
		status = inb(base + REG_STATUS);
	}
	if(status & STAT_ERR) {
		uint8_t err = inb(base + REG_ERR);
		panic("Error initializing ATA device: %X", err);
	}
	uint16_t data[256];
	for(int i = 0; i < 256; i++) {
		data[i] = inw(base + REG_DATA);
	}
	uint32_t sector_count = data[60] | (data[61] << 16);
	if(sector_count == 0) {
		panic("ATA device does not support LBA28");
	}
	dev->base = base;
	dev->ctrl = ctrl;
	dev->sector_count = sector_count;
	return true;
}

bool ata_init(struct AtaDevice *atadev, struct PciDevice pcidev) {
	uint8_t prog_if = pci_rcfg_b(pcidev, PCI_PROG_IF_B);

	if((prog_if & 0x7F) != 0) {
		panic("Could not initialize ATA device: unsupported prog_if %X", prog_if);
	}

	atadev->pcidev = pcidev;

	if(!ata_init_impl(0x1F0, 0x3F6, atadev)) {
		return false;
	}

	return true;
}

bool ata_read(struct AtaDevice atadev, uint32_t lba, uint8_t sector_count, uint16_t buf[sector_count*256]) {
	outb(atadev.base + REG_DRIVE_HEAD, 0xE0 | ((lba >> 24) & 0x0F));
	outb(atadev.base + REG_FEATURES, 0x00);
	outb(atadev.base + REG_SECTOR_CNT, sector_count);
	outb(atadev.base + REG_LBA_LO, lba & 0xFF);
	outb(atadev.base + REG_LBA_MID, (lba >> 8) & 0xFF);
	outb(atadev.base + REG_LBA_HI, (lba >> 16) & 0xFF);
	outb(atadev.base + REG_COMMAND, 0x20);
	for(int i = 0; i < sector_count; i++) {
		while(true) {
			uint8_t status = inb(atadev.ctrl);
			if(status & STAT_DRQ) break;
			if(status & STAT_DF) panic("Drive failure");
			if(status & STAT_ERR) panic("Read disk error");
		}
		for(int j = 0; j < 256; j++) {
			buf[i*256 + j] = inw(atadev.base + REG_DATA);
		}
	}
	return true;
}

bool ata_write(struct AtaDevice atadev, uint32_t lba, uint8_t sector_count, uint16_t buf[sector_count*256]) {
	outb(atadev.base + REG_DRIVE_HEAD, 0xE0 | ((lba >> 24) & 0x0F));
	outb(atadev.base + REG_FEATURES, 0x00);
	outb(atadev.base + REG_SECTOR_CNT, sector_count);
	outb(atadev.base + REG_LBA_LO, lba & 0xFF);
	outb(atadev.base + REG_LBA_MID, (lba >> 8) & 0xFF);
	outb(atadev.base + REG_LBA_HI, (lba >> 16) & 0xFF);
	outb(atadev.base + REG_COMMAND, 0x30);
	for(int i = 0; i < sector_count; i++) {
		while(true) {
			uint8_t status = inb(atadev.ctrl);
			if(status & STAT_DRQ) break;
			if(status & STAT_DF) panic("Drive failure");
			if(status & STAT_ERR) panic("Write disk error");
		}
		for(int j = 0; j < 256; j++) {
			outw(atadev.base + REG_DATA, buf[i*256 + j]);
		}
	}
	io_wait();
	outb(atadev.base + REG_COMMAND, 0xE7);
	while(inb(atadev.ctrl & STAT_BSY)) {}
	return true;
}