#include <stdio.h>

#include <port_io.h>

#include <drivers/ide.h>


static void io_delay() {
    for (int32_t i = 0; i < 1000; i++) {
        outb(0x80, 0);
    }
}


static int32_t ide_wait(uint16_t io, int32_t check_drq) {
    for (int32_t i = 0; i < 5000; i++) {
        uint8_t status = inb(io + ATA_REG_STATUS);
        if (!(status & ATA_SR_BSY)) {
            if (!check_drq || (status & ATA_SR_DRQ))
                return 0;
        }
        io_delay(1);
    }
    return 1;
}

static int32_t ide_wait_ready(uint16_t io) {
    for (int32_t i = 0; i < 100000; i++) {
        uint8_t status = inb(io + ATA_REG_STATUS);
        if (!(status & ATA_SR_BSY)) return 0;
    }
    return 1;
}



int32_t ide_identify(uint8_t drive, uint16_t* buffer) {
    uint16_t io = ATA_PRIMARY_IO;
    uint8_t slave_bit = (drive & 1) << 4;

    ide_wait(io, 0);
    outb(io + ATA_REG_HDDEVSEL, 0xA0 | slave_bit);
    outb(io + ATA_REG_SECCOUNT0, 0);
    outb(io + ATA_REG_LBA0, 0);
    outb(io + ATA_REG_LBA1, 0);
    outb(io + ATA_REG_LBA2, 0);
    outb(io + ATA_REG_COMMAND, 0xEC); // IDENTIFY DEVICE

    uint8_t status = inb(io + ATA_REG_STATUS);
    if (status == 0) return -1; // No device

    if (ide_wait(io, 1)) return -2;

    insw(io + ATA_REG_DATA, buffer, 256); // Read 512 bytes
    return 0;
}

void ide_initialize(void) {
    outb(ATA_PRIMARY_CTRL, 0x02); // Disable IRQs
    
    uint16_t identify_buf[256];
    if (ide_identify(0, identify_buf) == 0) {
        char model[41];
        for (int i = 0; i < 20; ++i) {
            model[i * 2 + 0] = identify_buf[27 + i] >> 8;
            model[i * 2 + 1] = identify_buf[27 + i] & 0xFF;
        }
        model[40] = 0;
        printf("Disk model: %s\n", model);
    }
}

int32_t ide_read48(uint8_t drive, uint64_t lba, uint8_t sector_count, void* buffer) {
    if (sector_count == 0) return -1;

    uint16_t io = ATA_PRIMARY_IO;
    uint8_t slave_bit = (drive & 1) << 4;

    ide_wait(io, 0); // Wait for BSY to clear

    outb(io + ATA_REG_HDDEVSEL, 0xE0 | slave_bit); // Select drive

    // High bytes (use 0xC2 - 0xC6 offsets if you define them separately)
    outb(io + 6, (sector_count >> 8) & 0xFF);       // Sector Count High
    outb(io + 7, (lba >> 40) & 0xFF);               // LBA[47:40]
    outb(io + 8, (lba >> 32) & 0xFF);               // LBA[39:32]
    outb(io + 9, (lba >> 24) & 0xFF);               // LBA[31:24]

    // Low bytes
    outb(io + 2, sector_count & 0xFF);              // Sector Count Low
    outb(io + 3, (lba >> 0) & 0xFF);                // LBA[7:0]
    outb(io + 4, (lba >> 8) & 0xFF);                // LBA[15:8]
    outb(io + 5, (lba >> 16) & 0xFF);               // LBA[23:16]


    outb(io + ATA_REG_COMMAND, ATA_CMD_READ_SECTORS_EXT);

    uint16_t* ptr = (uint16_t*)buffer;
    for (int i = 0; i < sector_count; i++) {
        if (ide_wait(io, 1)) {
            uint8_t status = inb(io + ATA_REG_STATUS);
            printf("IDE DRQ wait failed (read), status = 0x%02X\n", status);
            return -2;
        }

        uint8_t status = inb(io + ATA_REG_STATUS);
        if (status & ATA_SR_ERR) return -3;

        insw(io + ATA_REG_DATA, ptr, 256);
        ptr += 256;
    }

    return 0;
}


int32_t ide_write48(uint8_t drive, uint64_t lba, uint8_t sector_count, const void* buffer) {
    if (sector_count == 0) return -1;

    uint16_t io = ATA_PRIMARY_IO;
    uint8_t slave_bit = (drive & 1) << 4;

    ide_wait(io, 0); // Wait for BSY to clear

    outb(io + ATA_REG_HDDEVSEL, 0xE0 | slave_bit); // Select drive

    // Send high bytes of sector count and LBA first
    outb(io + ATA_REG_SECCOUNT0, 0);                     // Sector count high byte
    outb(io + ATA_REG_LBA0, (lba >> 40) & 0xFF);         // LBA[47:40]
    outb(io + ATA_REG_LBA1, (lba >> 32) & 0xFF);         // LBA[39:32]
    outb(io + ATA_REG_LBA2, (lba >> 24) & 0xFF);         // LBA[31:24]

    // Send low bytes
    outb(io + ATA_REG_SECCOUNT0, sector_count);          // Sector count low byte
    outb(io + ATA_REG_LBA0, (lba >> 0) & 0xFF);           // LBA[7:0]
    outb(io + ATA_REG_LBA1, (lba >> 8) & 0xFF);           // LBA[15:8]
    outb(io + ATA_REG_LBA2, (lba >> 16) & 0xFF);          // LBA[23:16]

    // Send write command
    outb(io + ATA_REG_COMMAND, ATA_CMD_WRITE_SECTORS_EXT); // 0x34

    const uint16_t* ptr = (const uint16_t*)buffer;
    for (int i = 0; i < sector_count; i++) {
        if (ide_wait(io, 1)) {
            uint8_t status = inb(io + ATA_REG_STATUS);
            printf("IDE DRQ wait failed (write), status = 0x%02X\n", status);
            return -2;
        }

        // Check for drive error
        uint8_t status = inb(io + ATA_REG_STATUS);
        if (status & ATA_SR_ERR) return -3;

        outsw(io + ATA_REG_DATA, ptr, 256); // 256 words = 512 bytes
        inb(io + ATA_REG_STATUS); /* Dummy read */
        ptr += 256;
    }

    // Wait for BSY=0 before issuing FLUSH CACHE
    if (ide_wait_ready(io)) return -5;
    outb(io + ATA_REG_COMMAND, 0xE7); // FLUSH CACHE
    if (ide_wait(io, 0)) {
      printf("FLUSH CACHE timeout, status=0x%02X\n", inb(io + ATA_REG_STATUS));
      return -4;
    }



    return 0;
}