#include <ata-pio.h>
#include <video.h>

uint16_t ata_identify[256];
// 0: disk is initialized
// 1: no disks
// 2: no disk on requested interface
// 3: ATAPI
int ata_init(uint8_t bus, uint8_t drive) {
	uint16_t io_base = ATA_IO_BASE - (bus ? ATA_SECONDARY_OFFSET : 0);
	if(inb(io_base + ATA_STATUS_REGISTER) == 0xFF) return 1;
	outb(io_base + ATA_DRIVE_REGISTER, drive ? 0xB0 : 0xA0);
	for(int i = 2; i < 6; i++) outb(io_base + i, 0);
	outb(io_base + ATA_STATUS_REGISTER, 0xEC);
	uint8_t status;
	do {
		status = inb(io_base + ATA_STATUS_REGISTER);
		if(!status) return 2;
	} while((status & 0x80));
	if(inb(io_base + 4) || inb(io_base + 5)) return 3;
	do {
		status = inb(io_base + ATA_STATUS_REGISTER);
		if(status & 1) return inb(io_base + ATA_ERROR_REGISTER);
	} while(!(status & 8));
	uint16_t err = inb(io_base + ATA_ERROR_REGISTER);
	if(err) return err;
	for(int i = 0; i < 256; i++) ata_identify[i] = inw(io_base);
	return 0;
}

uint16_t ata_setuprw(uint8_t bus, uint8_t drive, uint16_t sectors, uint16_t lba_high, uint32_t lba_low) {
	uint16_t io_base = ATA_IO_BASE - (bus ? ATA_SECONDARY_OFFSET : 0);
	outb(io_base + ATA_DRIVE_REGISTER, 0x40 | ((!!drive) << 4));
	uint8_t mports[4] = {2, 3, 4, 5};
	uint8_t dat[8] = {sectors >> 8, lba_low >> 24, lba_high >> 8, lba_high, sectors, lba_low, lba_low >> 8, lba_low >> 16};
	for(int i = 0; i < 8; i++) outb(io_base + mports[i % 4], dat[i]);
	return io_base;
}

int ata_read(uint8_t bus, uint8_t drive, uint16_t* buf, uint16_t sectors, uint16_t lba_high, uint32_t lba_low) {
	uint16_t io_base = ata_setuprw(bus, drive, sectors, lba_high, lba_low);
	outb(io_base + ATA_COMMAND_REGISTER, 0x24);
	for(int i = 0; i < sectors * 256; i++) buf[i] = inw(io_base);
}

int ata_write(uint8_t bus, uint8_t drive, uint16_t* buf, uint16_t sectors, uint16_t lba_high, uint32_t lba_low) {
	uint16_t io_base = ata_setuprw(bus, drive, sectors, lba_high, lba_low);
	outb(io_base + ATA_COMMAND_REGISTER, 0x34);
	for(int i = 0; i < sectors * 256; i++) outw(io_base, buf[i]);
	outb(io_base + ATA_COMMAND_REGISTER, 0xE7);
}