// ide.c - PIO-mode IDE driver (legacy 28-bit LBA)
#include <io_driver.h>
#include <hal_api.h>
#include <ke_debug.h>
#include <stdint.h>
#include <string.h>
#include <intrin.h>

// Primary channel I/O ports.
#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   // 0xE0=master 0xF0=slave
#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

#define CMD_READ   0x20
#define CMD_WRITE  0x30
#define CMD_IDENT  0xEC

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

static void IdeWait(void) {
    while (__inbyte(IDE_STATUS) & IDE_BUSY) _mm_pause();
}

static KSTATUS IdeWaitData(void) {
    // Read the Alternate Status port 4 times to waste 400ns
    for (int i = 0; i < 4; i++) {
        __inbyte(IDE_ALT);
    }

    // Wait until the drive is not busy AND data is requested (DRQ)
    while (1) {
        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;
        _mm_pause();
    }
}

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

    for (uint32_t i = 0; i < count; i++, lba++) {
        IdeWait(); // Wait for drive to be ready for commands
        __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,    CMD_READ);
        
        // Use the delay/DRQ wait
        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(uint64_t lba, uint32_t count, const void* buf) {
    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();
        __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,    CMD_WRITE);
        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;
    uint8_t status;
    int timeout;

    // 1. Select Primary Master Drive (0xA0)
    __outbyte(IDE_DRIVE, 0xA0);
    
    // Waste 400ns to let the drive selection settle on the bus
    for(int i=0; i<4; i++) __inbyte(IDE_ALT);

    // 2. Zero out the LBA and Sector Count registers
    __outbyte(IDE_SECCNT, 0);
    __outbyte(IDE_LBA0, 0);
    __outbyte(IDE_LBA1, 0);
    __outbyte(IDE_LBA2, 0);

    // 3. Send the IDENTIFY command
    DbgPrint("IDE: Sending CMD_IDENT (0xEC)...\n");
    __outbyte(IDE_CMD, CMD_IDENT);

    // 4. Read status immediately. If it's 0, the drive doesn't exist at all.
    status = __inbyte(IDE_STATUS);
    // DbgPrint("IDE: Initial status after IDENTIFY: 0x%02X\n", status);
    if (status == 0) {
        DbgPrint("IDE ERROR: Drive does not exist (Status 0).\n");
        return STATUS_NOT_FOUND;
    }

    // 5. Wait for BSY (Busy) to clear, WITH A TIMEOUT
    timeout = 100000;
    while (__inbyte(IDE_STATUS) & IDE_BUSY) {
        if (--timeout == 0) {
            DbgPrint("IDE ERROR: Timeout waiting for BSY to clear!\n");
            return STATUS_DEVICE_ERROR;
        }
    }

    // 6. Check for ATAPI (CD-ROM) or SATA signatures. 
    // If these registers aren't 0 after BSY clears, it's NOT a standard ATA disk!
    uint8_t lba1 = __inbyte(IDE_LBA1);
    uint8_t lba2 = __inbyte(IDE_LBA2);
    if (lba1 != 0 || lba2 != 0) {
        DbgPrint("IDE ERROR: Not an ATA drive! Signatures: LBA1=0x%02X, LBA2=0x%02X\n", lba1, lba2);
        // (ATAPI CD-ROMs usually show LBA1=0x14, LBA2=0xEB)
        return STATUS_NOT_SUPPORTED;
    }

    // 7. Wait for DRQ (Data Request) or ERR (Error)
    timeout = 100000;
    while (1) {
        status = __inbyte(IDE_STATUS);
        if (status & IDE_ERR_F) {
            DbgPrint("IDE ERROR: Drive returned an error! Status: 0x%02X\n", status);
            return STATUS_DEVICE_ERROR;
        }
        if (status & IDE_DRQ) {
            break; // Data is ready!
        }
        if (--timeout == 0) {
            DbgPrint("IDE ERROR: Timeout waiting for DRQ! Status: 0x%02X\n", status);
            return STATUS_DEVICE_ERROR;
        }
    }

    // 8. Success! Read the 256 words (512 bytes) of IDENTIFY data
    // DbgPrint("IDE: DRQ set! Reading 256 words...\n");
    uint16_t ident[256];
    for (int i = 0; i < 256; i++) {
        ident[i] = __inword(IDE_DATA);
    }

    // 9. Extract and print the sector count (Words 60 and 61 for 28-bit LBA)
    gSectorCount = ((uint32_t)ident[61] << 16) | ident[60];
    DbgPrint("IDE: SUCCESS: Disk found! Sector count: %u\n", (uint32_t)gSectorCount);

    if (gSectorCount == 0) return STATUS_NOT_FOUND;

    // Register driver globally (from the previous fix!)
    gDrv.ReadSectors  = IdeRead;
    gDrv.WriteSectors = IdeWrite;
    gDrv.SectorCount  = gSectorCount;

    return IoRegisterDisk(0, &gDrv);
}