#include <kernel/ide.h>
#include <kernel/io.h>
#include <kernel/printf.h>

uint16_t ide_select_drive(uint8_t bus, uint8_t drive) {
    if (bus == ATA_PRIMARY) {
        if (drive == ATA_MASTER)
            outportb(ATA_PRIMARY_IO + ATA_REG_HDDEVSEL, 0xA0);
        else
            outportb(ATA_PRIMARY_IO + ATA_REG_HDDEVSEL, 0xB0);

        return ATA_PRIMARY_IO;
    } else {
        if (drive == ATA_MASTER)
            outportb(ATA_SECONDARY_IO + ATA_REG_HDDEVSEL, 0xA0);
        else
            outportb(ATA_SECONDARY_IO + ATA_REG_HDDEVSEL, 0xB0);

        return ATA_SECONDARY_IO;
    }
}

uint8_t ide_identify(uint8_t bus, uint8_t drive, uint16_t *ide_buf) {
    asm volatile ("cli" ::);

    uint16_t io = ide_select_drive(bus, drive);
    // Must zero these registers before sending IDENTIFY command
    outportb(io + ATA_REG_SECCOUNT0, 0);
	outportb(io + ATA_REG_LBA0, 0);
	outportb(io + ATA_REG_LBA1, 0);
	outportb(io + ATA_REG_LBA2, 0);

    for (int i = 0; i < 14; i++)
        inportb(io + ATA_REG_STATUS);
    uint8_t status = inportb(io + ATA_REG_STATUS);

    // Send IDENTIFY command
    outportb(io + ATA_REG_COMMAND, ATA_CMD_IDENTIFY);

    status = inportb(io + ATA_REG_STATUS);
    // Make sure there are no errors
    if (!status) {
        printf("Drive does not exist\n");
        asm volatile ("sti" ::);
        return 0;
    }

    // Poll the drive until it's not busy
    while (inportb(io + ATA_REG_STATUS) & ATA_SR_BSY) 
        asm volatile ("pause" ::);

    if (inportb(io + ATA_REG_LBA1) || inportb(io + ATA_REG_LBA1)) {
        printf("Drive is not ATA. wtf?");
        asm volatile ("sti" ::);
        return 0;
    }

    // Poll until drive is ready for data requests or until an error occurs
    status = inportb(io + ATA_REG_STATUS);
    while (!(status & ATA_SR_DRQ) && !(status & ATA_SR_ERR)) {
        status = inportb(io + ATA_REG_STATUS);
    }

    if (status & ATA_SR_ERR) {
        printf("DRIVE ERROR\n");
        asm volatile ("sti" ::);
        return 0;
    }

    printf("%s %s is online. Status: %X\n", bus==ATA_PRIMARY?"Primary":"Secondary", drive==ATA_PRIMARY?"master":"slave", status);
    // Read identification space of the device
    for (int i = 0; i < 256; i++) {
       *(ide_buf + i*2) = inportw(io + ATA_REG_DATA);
    }

    asm volatile ("sti" ::);
    return 0;
}

uint8_t ata_read_sector(uint8_t bus, uint8_t drive, uint32_t lba, uint8_t *buf) {
    asm volatile ("cli" ::);
    uint16_t io = 0;
    if (bus == ATA_PRIMARY) {
        io = ATA_PRIMARY_IO;
    } else {
        io = ATA_SECONDARY_IO;
    }

	uint8_t cmd = (drive==ATA_MASTER?0xE0:0xF0);
    outportb(io + ATA_REG_HDDEVSEL, cmd | (uint8_t)(lba >> 24 & 0x0f));
    outportb(io + 1, 0x00);
    outportb(io + ATA_REG_SECCOUNT0, 1);
    outportb(io + ATA_REG_LBA0, (uint8_t)((lba)));
	outportb(io + ATA_REG_LBA1, (uint8_t)((lba) >> 8));
	outportb(io + ATA_REG_LBA2, (uint8_t)((lba) >> 16));
	outportb(io + ATA_REG_COMMAND, ATA_CMD_READ_PIO);


    uint8_t status = inportb(io + ATA_REG_STATUS);
    while (!(status & ATA_SR_DRQ) && !(status & ATA_SR_ERR)) {
        status = inportb(io + ATA_REG_STATUS);
    }

    for (int i = 0; i < 256; i++)
        *(uint16_t *)(buf + i*2) = inportw(io + ATA_REG_DATA);

    asm volatile ("sti" ::);
    return 1;
}