#include <stdint.h>
#include <string.h>
#include <intrin.h>
#include <io_driver.h>
#include <hal_api.h>
#include <ke_debug.h>
#include <ke_sched.h>

#define IDE_DATA    0x1F0
#define IDE_ERR     0x1F1
#define IDE_SECCNT  0x1F2
#define IDE_LBA0    0x1F3
#define IDE_LBA1    0x1F4
#define IDE_LBA2    0x1F5
#define IDE_DRIVE   0x1F6   
#define IDE_STATUS  0x1F7
#define IDE_CMD     0x1F7
#define IDE_ALT     0x3F6

#define IDE_BUSY   0x80
#define IDE_READY  0x40
#define IDE_ERR_F  0x01
#define IDE_DRQ    0x08

// 48-bit LBA Commands
#define CMD_READ_EXT   0x24
#define CMD_WRITE_EXT  0x34
#define CMD_IDENT      0xEC

static const HAL_API* gHal;
static DISK_DRIVER gDrv;
static uint64_t gSectorCount = 0;
static bool gSupportsLba48 = false;

// ---------------------------------------------------------
// Polling Upgrades: Yield instead of blocking the CPU
// ---------------------------------------------------------
static void IdeWait(void) {
    while (__inbyte(IDE_STATUS) & IDE_BUSY) {
        KeYield(); // <-- Let other threads run while the disk thinks!
    }
}

static KSTATUS IdeWaitData(void) {
    for (int i = 0; i < 4; i++) __inbyte(IDE_ALT);

    while (true) {
        uint8_t st = __inbyte(IDE_STATUS);
        if (st & IDE_ERR_F) return STATUS_DEVICE_ERROR;
        if (!(st & IDE_BUSY) && (st & IDE_DRQ)) return STATUS_SUCCESS;
        KeYield(); // <-- Be a good citizen
    }
}

// ---------------------------------------------------------
// 48-Bit LBA Read/Write
// ---------------------------------------------------------
static KSTATUS IdeRead(void *ctx, uint64_t lba, uint32_t count, void* buf) {
    (void)ctx;
    if (lba + count > gSectorCount) return STATUS_INVALID_PARAMETER;
    uint16_t* p = (uint16_t *)buf;

    // For simplicity, handle one sector at a time in this loop
    for (uint32_t i = 0; i < count; i++, lba++) {
        IdeWait(); 

        if (gSupportsLba48) {
            // LBA48 requires writing high bytes first, then low bytes
            __outbyte(IDE_DRIVE,  0x40); // 0x40 = LBA mode, drive 0
            __outbyte(IDE_SECCNT, 0);                 // Count High
            __outbyte(IDE_LBA0,   (lba >> 24) & 0xFF); // LBA 4
            __outbyte(IDE_LBA1,   (lba >> 32) & 0xFF); // LBA 5
            __outbyte(IDE_LBA2,   (lba >> 40) & 0xFF); // LBA 6
            
            __outbyte(IDE_SECCNT, 1);                 // Count Low
            __outbyte(IDE_LBA0,   (uint8_t)lba);       // LBA 1
            __outbyte(IDE_LBA1,   (uint8_t)(lba >> 8)); // LBA 2
            __outbyte(IDE_LBA2,   (uint8_t)(lba >> 16));// LBA 3
            
            __outbyte(IDE_CMD,    CMD_READ_EXT);
        } else {
            // Fallback to 28-bit
            __outbyte(IDE_DRIVE,  0xE0 | ((lba >> 24) & 0xF));
            __outbyte(IDE_SECCNT, 1);
            __outbyte(IDE_LBA0,   (uint8_t)lba);
            __outbyte(IDE_LBA1,   (uint8_t)(lba >> 8));
            __outbyte(IDE_LBA2,   (uint8_t)(lba >> 16));
            __outbyte(IDE_CMD,    0x20); // CMD_READ (28-bit)
        }
        
        if (IdeWaitData() != STATUS_SUCCESS) return STATUS_DEVICE_ERROR;
        for (int w = 0; w < 256; w++) *p++ = __inword(IDE_DATA);
    }
    return STATUS_SUCCESS;
}

static KSTATUS IdeWrite(void *ctx, uint64_t lba, uint32_t count, const void* buf) {
    (void)ctx;
    if (lba + count > gSectorCount) return STATUS_INVALID_PARAMETER;
    const uint16_t* p = (const uint16_t *)buf;

    for (uint32_t i = 0; i < count; i++, lba++) {
        IdeWait();
        
        if (gSupportsLba48) {
            __outbyte(IDE_DRIVE,  0x40);
            __outbyte(IDE_SECCNT, 0);
            __outbyte(IDE_LBA0,   (lba >> 24) & 0xFF);
            __outbyte(IDE_LBA1,   (lba >> 32) & 0xFF);
            __outbyte(IDE_LBA2,   (lba >> 40) & 0xFF);
            
            __outbyte(IDE_SECCNT, 1);
            __outbyte(IDE_LBA0,   (uint8_t)lba);
            __outbyte(IDE_LBA1,   (uint8_t)(lba >> 8));
            __outbyte(IDE_LBA2,   (uint8_t)(lba >> 16));
            
            __outbyte(IDE_CMD,    CMD_WRITE_EXT);
        } else {
            __outbyte(IDE_DRIVE,  0xE0 | ((lba >> 24) & 0xF));
            __outbyte(IDE_SECCNT, 1);
            __outbyte(IDE_LBA0,   (uint8_t)lba);
            __outbyte(IDE_LBA1,   (uint8_t)(lba >> 8));
            __outbyte(IDE_LBA2,   (uint8_t)(lba >> 16));
            __outbyte(IDE_CMD,    0x30); // CMD_WRITE (28-bit)
        }

        IdeWait();
        for (int w = 0; w < 256; w++) __outword(IDE_DATA, *p++);
        __outbyte(IDE_ALT, 0); // flush cache
        
        if (IdeWaitData() != STATUS_SUCCESS) return STATUS_DEVICE_ERROR;
    }
    return STATUS_SUCCESS;
}

KSTATUS DriverEntry(const HAL_API* hal) {
    DbgPrint("IDE: Init\n");
    gHal = hal;

    // Select Primary Master (Drive 0)
    __outbyte(IDE_DRIVE, 0xA0);
    for(int i=0; i<4; i++) __inbyte(IDE_ALT);

    __outbyte(IDE_SECCNT, 0);
    __outbyte(IDE_LBA0, 0);
    __outbyte(IDE_LBA1, 0);
    __outbyte(IDE_LBA2, 0);

    __outbyte(IDE_CMD, CMD_IDENT);
    
    if (__inbyte(IDE_STATUS) == 0) {
        // QEMU usually falls in here because it uses AHCI by default, 
        // leaving the legacy IDE ports floating/unconnected.
        DbgPrint("IDE ERROR: Drive does not exist on legacy port 0x1F0.\n");
        return STATUS_NOT_FOUND;
    }

    // Wait for BSY to clear, yielding to scheduler
    int timeout = 10000;
    while (__inbyte(IDE_STATUS) & IDE_BUSY) {
        if (--timeout == 0) return STATUS_DEVICE_ERROR;
        KeYield(); 
    }

    if (__inbyte(IDE_LBA1) != 0 || __inbyte(IDE_LBA2) != 0) {
        DbgPrint("IDE ERROR: Not a standard ATA drive (Possibly ATAPI).\n");
        return STATUS_NOT_SUPPORTED;
    }

    timeout = 10000;
    while (true) {
        uint8_t status = __inbyte(IDE_STATUS);
        if (status & IDE_ERR_F) return STATUS_DEVICE_ERROR;
        if (status & IDE_DRQ) break;
        if (--timeout == 0) return STATUS_DEVICE_ERROR;
        KeYield();
    }

    uint16_t ident[256];
    for (int i = 0; i < 256; i++) {
        ident[i] = __inword(IDE_DATA);
    }

    // Word 83 Bit 10 determines LBA48 support
    gSupportsLba48 = (ident[83] & (1 << 10)) != 0;

    if (gSupportsLba48) {
        // Words 100-103 contain the 64-bit LBA sector count
        gSectorCount = *(uint64_t*)&ident[100];
        DbgPrint("IDE: LBA48 Disk found! Sectors: %llu\n", gSectorCount);
    } else {
        // Fallback to 28-bit
        gSectorCount = ((uint32_t)ident[61] << 16) | ident[60];
        DbgPrint("IDE: LBA28 Disk found! Sectors: %u\n", (uint32_t)gSectorCount);
    }

    if (gSectorCount == 0) return STATUS_NOT_FOUND;

    gDrv.Ctx          = NULL;
    gDrv.ReadSectors  = IdeRead;
    gDrv.WriteSectors = IdeWrite;
    gDrv.SectorCount  = gSectorCount;

    return IoRegisterDisk(0, &gDrv);
}