#include "drivers/atapi/ata.h"

#include "lib/print.h"

#include <stddef.h>
#include <stdint.h>

extern uint8_t inb(uint16_t port);
extern void outb(uint16_t port, uint8_t data);

extern void insw(uint32_t address, void* buf, uint32_t sz);
extern void outsw(uint32_t address, void* buf, uint32_t sz);

uint8_t selected_drive_primary = 0;
uint8_t selected_drive_secondary = 0;

uint8_t ide_info_buffer[512];

ide_device_t ide_devices[4];

static void fix_ata_string(char* buf, size_t sz) {
    if ((sz % 2) != 0) return;
    for (size_t i = 0; i < sz; i += 2) {
        char tmp = buf[i];
        buf[i] = buf[i + 1];
        buf[i + 1] = tmp;
    }
    buf[sz-1] = '\0'; 
}

static uint16_t fetch_ctrl_reg(uint16_t bus) {
    if (bus == IDE_PRIMARY_IO_BASE) return IDE_PRIMARY_CRTL_BASE;
    if (bus == IDE_SECONDARY_IO_BASE) return IDE_SECONDARY_CRTL_BASE;
    return 0;
}   

static void ata_delay(uint16_t bus) {
    uint16_t ctrl = fetch_ctrl_reg(bus);
    for (int i = 0; i < 5; i++) {
        inb(ctrl);
    }
}

static int wait_not_bsy_drq(uint16_t bus) {
    ata_delay(bus);

    int timeout = 10000;
    int error = 0;
    while(timeout--) {
        uint16_t ctrl = fetch_ctrl_reg(bus);
        uint8_t stat = inb(ctrl);
        if (stat & STATUS_REG_ERR) error = 1; break;
        if (!(stat & STATUS_REG_BSY)) break;
    }

    if (timeout == 0 || error) {
        return 1;
    }

    return 0;
}

static int ide_irq_invoked = 0;
static void ide_irq_wait() {
    while (!ide_irq_invoked)
      ;
   ide_irq_invoked = 0;
}

void ide_invoke_irq() {
    ide_irq_invoked = 1;
}

int select_drive(uint16_t bus, uint8_t drive) {
    if (bus == IDE_PRIMARY_IO_BASE && selected_drive_primary == drive) {
        return 1;
    }

    if (bus == IDE_SECONDARY_IO_BASE && selected_drive_secondary == drive) {
        return 1;
    }

    uint16_t ctrl_reg = fetch_ctrl_reg(bus);
    if (ctrl_reg == 0) {
        return 0;
    }

    // wait for bsy and drq clear
    while (1) {
        uint8_t status = inb(ctrl_reg + CTRL_ALT_STATUS_REG);
        if (!(status & STATUS_REG_BSY) && !(status & STATUS_REG_DRQ)) {
            break;
        }
    }

    outb(bus + IO_DRIVE_SELECT_REG, drive);

    ata_delay(bus);

    uint8_t status = inb(ctrl_reg + CTRL_ALT_STATUS_REG);
    if (status & STATUS_REG_ERR) {
        return 0;
    }

    return 1;
}

void ident_device(uint16_t bus, uint8_t drive, uint32_t slot) {
    ide_dev_type type = NO_DEVICE;

    uint8_t stat = inb(bus + IO_STATUS_REG);
    if (stat == 0xFF) {
        kprintf("no device on bus 0x%x\n", bus);
        return;
    }

    if (!select_drive(bus, drive)) {
        kprintf("failed to select drive 0x%x on bus 0x%x\n", bus, drive);
        return;
    }

    outb(bus + IO_SECTOR_COUNT_REG, 0x00);
    outb(bus + IO_LBA_LO_REG, 0x00);
    outb(bus + IO_LBA_MID_REG, 0x00);
    outb(bus + IO_LBA_HI_REG, 0x00);
    // ident ata
    outb(bus + IO_CMD_REG, ATA_CMD_IDENTIFY);
    ata_delay(bus);

    stat = inb(bus + IO_STATUS_REG);
    if (stat == 0x00) {
        kprintf("no drive 0x%x on bus 0x%x\n", drive, bus);
        return;
    }

    int not_ata = 0;
    int timeout = 10000;
    while(timeout--) {
        stat = inb(bus + IO_STATUS_REG);
        if (stat & STATUS_REG_ERR) not_ata = 1; break;
        if (!(stat & STATUS_REG_BSY) && (stat & STATUS_REG_DRQ)) type = ATA; break;
    }

    if (timeout == 0) {
        kprintf("drive 0x%x ident timed out on bus 0x%x\n", drive, bus);
        return;
    }

    if (not_ata) {
        uint8_t mid = inb(bus + IO_LBA_MID_REG);
        uint8_t hi = inb(bus + IO_LBA_HI_REG);

        if (mid != 0x14 || hi != 0xEB) {
            return;
        }

        type = ATAPI;

        // ident atapi
        outb(bus + IO_CMD_REG, ATA_CMD_IDENTIFY_PACKET);
        ata_delay(bus);
    }

    
    insw(bus + IO_DATA_REG, ide_info_buffer, 256);

    ide_devices[slot].bus = bus;
    ide_devices[slot].device = drive;
    ide_devices[slot].type = type;
    ide_devices[slot].capabilities = *(uint16_t*)(ide_info_buffer + ATA_IDENT_CAPABILITIES);
    ide_devices[slot].sig = *(uint16_t*)(ide_info_buffer + ATA_IDENT_DEVICETYPE);
    ide_devices[slot].command_set = *(uint32_t*)(ide_info_buffer + ATA_IDENT_COMMANDSETS);

    if (ide_devices[slot].command_set & (1 << 26)) {
        ide_devices[slot].size = *(uint32_t*)(ide_info_buffer + ATA_IDENT_MAX_LBA_EXT);
    } else {
        ide_devices[slot].size = *(uint32_t*)(ide_info_buffer + ATA_IDENT_MAX_LBA);
    }

    char* serial_buf = (char*)(ide_info_buffer + ATA_IDENT_SERIAL);
    fix_ata_string(serial_buf, 8);
    // ide_devices[slot].serial = serial_buf;

    char* model_buf = (char*)(ide_info_buffer + ATA_IDENT_MODEL);
    fix_ata_string(model_buf, ATA_IDENT_MODEL_SZ);
    // ide_devices[slot].model = model_buf;

    kprintf("Found device 0x%x on bus 0x%x:\n", drive, bus);
    kprintf("\tModel Number: %s\n", model_buf);
    kprintf("\tSerial Number: %s\n", serial_buf);
    kprintf("\tSize %dB\n", *(uint32_t*)(ide_info_buffer + ATA_IDENT_MAX_LBA));
}

uint8_t temp_sector_buf[2048];
void atapi_read_sector(uint8_t drive_num, uint32_t lba, uint32_t sectors, void* buf) {
    ide_device_t drive = ide_devices[drive_num];

    // enable irqs
    uint16_t ctrl_reg = fetch_ctrl_reg(drive.bus);
    outb(ctrl_reg, 0x00);

    uint8_t pkt[12];
    pkt[ 0] = ATAPI_CMD_READ;
    pkt[ 1] = 0x0;
    pkt[ 2] = (lba >> 24) & 0xFF;
    pkt[ 3] = (lba >> 16) & 0xFF;
    pkt[ 4] = (lba >> 8) & 0xFF;
    pkt[ 5] = (lba >> 0) & 0xFF;
    pkt[ 6] = 0x0;
    pkt[ 7] = 0x0;
    pkt[ 8] = 0x0;
    pkt[ 9] = sectors;
    pkt[10] = 0x0;
    pkt[11] = 0x0;

    if (!select_drive(drive.bus, drive.device)) {
        return;
    }

    outb(drive.bus + IO_FEATURE_REG, 0x00); // PIO mode

    uint32_t words = 1024;
    outb(drive.bus + IO_LBA_MID_REG, (words * 2) & 0xFF);
    outb(drive.bus + IO_LBA_HI_REG, (words * 2) >> 8);

    outb(drive.bus + IO_CMD_REG, ATA_CMD_PACKET);
    // poll for drq
    int err = wait_not_bsy_drq(drive.bus);
    if (err) {
        kprintf("failed sending ata packet command\n");
        return;
    }
    outsw(drive.bus + IO_DATA_REG, pkt, 6);

    uint32_t ptr = (uint32_t)buf;
    for (int i = 0; i < sectors; i++) {
        ide_irq_wait();
        err = wait_not_bsy_drq(drive.bus);
        if (err) {
            kprintf("failed reading atapi sector\n");
            return;
        }
        
        insw(drive.bus + IO_DATA_REG, (void*)ptr, 1024);
        ptr += 2048;    
    }

    ide_irq_wait();
    while(inb(drive.bus + IO_STATUS_REG) & (STATUS_REG_BSY | STATUS_REG_DRQ));

    kprintf("read %d sectors from drive %d\n", sectors, drive_num);
}

void scan_drives() {
    ident_device(IDE_SECONDARY_IO_BASE, IDE_DRIVE_MASTER, 0);
    
    ide_device_t boot_drive = ide_devices[0];

    if (!(boot_drive.capabilities & 0x200)) {
        kprintf("boot drive does not support LBA addressing");
        return;
    }

    atapi_read_sector(0, 16, 1, temp_sector_buf);
}