// ide.c - IDE/ATA disk I/O path and DriverEntry.
#include "ide.h"

// ---------------------------------------------------------------------------
// Module state
// ---------------------------------------------------------------------------
const HAL_API *gHal;
DRIVER_OBJECT  gDriver;
DEVICE_OBJECT *gDevice = NULL;
bool           gDmaEnabled = false;
uint64_t       gSectorCount   = 0;
bool           gSupportsLba48 = false;

// DMA hardware state
uint16_t  gBmBase   = 0;     // Bus Master I/O base port (BAR4 & ~3)
PRD      *gPrdtVirt = NULL;  // kernel VA of PRDT page
uint32_t  gPrdtPhys = 0;    // 32-bit physical address of PRDT

// Async worker synchronisation
IRP_QUEUE  gIrpQueue;
SEMAPHORE *gIrpSemaphore = NULL;
EVENT     *gDmaEvent     = NULL;  // auto-reset; DPC signals after each xfer
KDPC       gDmaDpc;               // bottom half for DMA completion

// ---------------------------------------------------------------------------
// Bus Master I/O helpers
// ---------------------------------------------------------------------------
static inline uint8_t  BmRead8 (uint8_t reg)           { return __inbyte (gBmBase + reg); }
static inline void     BmWrite8 (uint8_t reg, uint8_t v)  { __outbyte (gBmBase + reg, v); }
static inline void     BmWrite32(uint8_t reg, uint32_t v) { __outdword(gBmBase + reg, v); }

// ---------------------------------------------------------------------------
// PIO polling helpers (unchanged)
// ---------------------------------------------------------------------------
static void IdeWait(void) {
    while (__inbyte(IDE_STATUS) & IDE_BUSY)
        KeYield();
}

static KSTATUS IdeWaitData(void) {
    for (int i = 0; i < 4; i++) __inbyte(IDE_ALT);  // 400 ns delay
    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();
    }
}

// ---------------------------------------------------------------------------
// PIO read/write — sync fallback when DMA is unavailable
// ---------------------------------------------------------------------------
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 (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_READ_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,    0x20);
        }
        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);
        }
        IdeWait();
        for (int w = 0; w < 256; w++) __outword(IDE_DATA, *p++);
        __outbyte(IDE_ALT, 0);   // flush write cache
        if (IdeWaitData() != STATUS_SUCCESS) return STATUS_DEVICE_ERROR;
    }
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// BuildPrdt — fill gPrdtVirt[] for a transfer.
//
// Walks the virtual buffer in page-aligned chunks, looks up each chunk's
// physical address via MmGetPhysicalAddress, and coalesces adjacent physical
// pages into a single PRD entry where possible. This is correct regardless of
// whether the pool gave us physically contiguous pages.
//
// Constraints enforced:
//   - Each PRD entry covers at most one 4 KB physical page (never crosses
//     a page boundary), so the physical address is valid for the whole chunk.
//   - Total byte count per entry is ≤ 64 KB (the PRD limit).
//   - Physical address must fit in 32 bits (BM IDE hardware limit).
//
// Returns the number of PRD entries written, or 0 on error.
// ---------------------------------------------------------------------------
static uint32_t BuildPrdt(void *virtBuf, uint32_t byteLen) {
    uint32_t  nPrd   = 0;
    uintptr_t va     = (uintptr_t)virtBuf;
    uint32_t  remain = byteLen;

    while (remain > 0) {
        if (nPrd >= PRDT_MAX_ENTRIES) {
            DbgPrint("IDE DMA: transfer too large for PRDT (%u entries)\n", nPrd);
            return 0;
        }

        // Bytes from va to the end of its 4 KB page (never straddles a page).
        uint32_t pageOff  = (uint32_t)(va & 0xFFF);
        uint32_t inPage   = 0x1000 - pageOff;
        uint32_t chunk    = (remain < inPage) ? remain : inPage;

        PHYSICAL_ADDRESS phys = MmGetPhysicalAddress(va);
        if (phys >> 32) {
            // Physical address above 4 GB — BM IDE's 32-bit PRDT can't reach it.
            DbgPrint("IDE DMA: buffer page at VA 0x%016llX has phys 0x%016llX > 4 GB\n",
                     (uint64_t)va, (uint64_t)phys);
            return 0;
        }
        uint32_t phys32 = (uint32_t)phys;

        // Try to coalesce with the previous PRD entry if the physical regions
        // are adjacent and the combined size stays within the 64 KB limit.
        if (nPrd > 0) {
            PRD     *prev      = &gPrdtVirt[nPrd - 1];
            uint32_t prevSz    = prev->byteCount ? prev->byteCount : 0x10000u;
            uint32_t prevEnd   = prev->physAddr + prevSz;
            uint32_t combined  = prevSz + chunk;
            if (prevEnd == phys32 && combined <= 0x10000u) {
                prev->byteCount = (uint16_t)(combined == 0x10000u ? 0u : combined);
                va     += chunk;
                remain -= chunk;
                continue;
            }
        }

        gPrdtVirt[nPrd].physAddr  = phys32;
        gPrdtVirt[nPrd].byteCount = (uint16_t)(chunk == 0x10000u ? 0u : chunk);
        gPrdtVirt[nPrd].flags     = 0;
        nPrd++;
        va     += chunk;
        remain -= chunk;
    }

    if (nPrd > 0)
        gPrdtVirt[nPrd - 1].flags |= PRD_EOT;

    return nPrd;
}

// ---------------------------------------------------------------------------
// IDE DMA completion DPC.  The ISR owns hardware acknowledgement; this bottom
// half wakes the worker outside hard IRQ/DIRQL context.
// ---------------------------------------------------------------------------
void IdeDmaCompletionDpc(KDPC *dpc, void *context,
                                void *systemArgument1, void *systemArgument2) {
    (void)dpc;
    (void)context;
    (void)systemArgument1;
    (void)systemArgument2;

    if (gDmaEvent)
        KeSignalEvent(gDmaEvent);
}

// ---------------------------------------------------------------------------
// IDE ISR (IRQ 14 — primary channel)
//
// Must be fast — runs with interrupts disabled. We only clear the hardware
// interrupt bits here and queue a DPC to wake the worker thread.
// The worker reads the final status and fires IRP completion hooks.
// ---------------------------------------------------------------------------
void IdeIsr(void) {
    uint8_t bmSt = BmRead8(BM_STATUS_REG);
    if (!(bmSt & BM_STATUS_IRQ))
        return;   // not our interrupt (spurious or secondary channel)

    // Stop the Bus Master engine.
    BmWrite8(BM_CMD_REG, 0x00);

    // Clear interrupt and error sticky bits (write-1-to-clear).
    BmWrite8(BM_STATUS_REG, BM_STATUS_IRQ | BM_STATUS_ERROR);

    // Read IDE status and alt-status to fully acknowledge the device interrupt.
    (void)__inbyte(IDE_ALT);
    (void)__inbyte(IDE_STATUS);

    // Wake happens in the DPC. gDmaEvent is auto-reset so exactly one waiter
    // is released when the bottom half runs.
    KeInsertQueueDpc(&gDmaDpc, (void *)0, (void *)0);
}

// ---------------------------------------------------------------------------
// IdeDmaRw — program the hardware and block until the ISR completes it.
//
// Called exclusively from IdeWorker. Has full control of the IDE/BM
// registers; no other thread touches them while the worker is running.
// Falls back to PIO inline if BuildPrdt fails (address space issue).
// ---------------------------------------------------------------------------
static KSTATUS IdeDmaRw(uint64_t lba, uint32_t count, void *buf, bool write) {
    uint32_t byteLen = count * 512u;

    uint32_t nPrd = BuildPrdt(buf, byteLen);
    if (nPrd == 0) {
        // Graceful inline PIO fallback for awkward buffers.
        DbgPrint("IDE DMA: BuildPrdt failed — falling back to PIO\n");
        return write ? IdeWrite(NULL, lba, count, buf)
                     : IdeRead (NULL, lba, count, buf);
    }

    // ---- Program Bus Master -----------------------------------------------

    // Load PRDT physical address.
    BmWrite32(BM_PRDT_REG, gPrdtPhys);

    // Clear any leftover interrupt/error bits from the previous transfer.
    BmWrite8(BM_STATUS_REG, BM_STATUS_IRQ | BM_STATUS_ERROR);

    // Set direction: BM_CMD_READ = disk→memory; 0 = memory→disk.
    // Do NOT set BM_CMD_START yet — spec says program ATA regs first.
    BmWrite8(BM_CMD_REG, write ? 0x00 : BM_CMD_READ);

    // ---- Program ATA registers -------------------------------------------
    IdeWait();

    if (gSupportsLba48) {
        __outbyte(IDE_DRIVE,  0x40);
        __outbyte(IDE_SECCNT, (uint8_t)(count >> 8));   // sector count high byte
        __outbyte(IDE_LBA0,   (uint8_t)(lba >> 24));
        __outbyte(IDE_LBA1,   (uint8_t)(lba >> 32));
        __outbyte(IDE_LBA2,   (uint8_t)(lba >> 40));
        __outbyte(IDE_SECCNT, (uint8_t)count);           // sector count low byte
        __outbyte(IDE_LBA0,   (uint8_t)lba);
        __outbyte(IDE_LBA1,   (uint8_t)(lba >>  8));
        __outbyte(IDE_LBA2,   (uint8_t)(lba >> 16));
        __outbyte(IDE_CMD,    write ? CMD_WRITE_DMA_EXT : CMD_READ_DMA_EXT);
    } else {
        __outbyte(IDE_DRIVE,  0xE0 | ((lba >> 24) & 0xF));
        __outbyte(IDE_SECCNT, (uint8_t)count);
        __outbyte(IDE_LBA0,   (uint8_t)lba);
        __outbyte(IDE_LBA1,   (uint8_t)(lba >>  8));
        __outbyte(IDE_LBA2,   (uint8_t)(lba >> 16));
        __outbyte(IDE_CMD,    write ? CMD_WRITE_DMA : CMD_READ_DMA);
    }

    // ---- Start DMA --------------------------------------------------------
    // Set Start bit. Direction bit is already set from above.
    BmWrite8(BM_CMD_REG, (write ? 0x00 : BM_CMD_READ) | BM_CMD_START);

    // Block the worker thread until the ISR fires.
    KeWaitForEvent(gDmaEvent, 0);

    // ---- Read result -------------------------------------------------------
    // ISR already stopped the BM engine and cleared BM_STATUS_IRQ.
    // We still need to check BM_STATUS_ERROR and IDE_ERR_F.
    uint8_t bmSt  = BmRead8(BM_STATUS_REG);
    uint8_t ideSt = __inbyte(IDE_STATUS);

    if (bmSt & BM_STATUS_ERROR) {
        DbgPrint("IDE DMA: BM error (bmStatus=0x%02x, ideStatus=0x%02x, ideErr=0x%02x)\n",
                 bmSt, ideSt, __inbyte(IDE_ERR));
        return STATUS_DEVICE_ERROR;
    }
    if (ideSt & IDE_ERR_F) {
        DbgPrint("IDE DMA: ATA error (ideStatus=0x%02x, ideErr=0x%02x)\n",
                 ideSt, __inbyte(IDE_ERR));
        return STATUS_DEVICE_ERROR;
    }
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// IdeWorker — kernel thread, one per IDE controller.
//
// Blocks on gIrpSemaphore. Each IrpQueueEnqueue increments it by one, so
// the worker wakes once per submitted IRP. Mirrors AhciWorker exactly.
// ---------------------------------------------------------------------------
void IdeWorker(uintptr_t arg) {
    (void)arg;
    while (true) {
        KeWaitSemaphore(gIrpSemaphore);

        IRP *irp = IrpQueueDequeue(&gIrpQueue);
        if (!irp)
            KePanic(KE_PANIC_CODE_DRIVER, "IdeWorker: semaphore fired but IRP queue is empty — accounting bug\n");

        KSTATUS st;
        if (irp->Lba + irp->SectorCount > gSectorCount) {
            st = STATUS_INVALID_PARAMETER;
        } else {
            bool write = !!(irp->Flags & IRP_FLAG_WRITE);
            st = IdeDmaRw(irp->Lba, irp->SectorCount, irp->Buffer, write);
        }

        IoCompleteRequest(irp, st, KSUCCESS(st) ? ((size_t)irp->SectorCount * 512u) : 0);
    }
}

// ---------------------------------------------------------------------------
// IdeDispatchReadWrite — DEVICE_OBJECT dispatch entry for disk IRPs.
//
// If DMA is online, the IRP is queued to the worker and completed later.
// If DMA is unavailable, the IRP is completed inline through the PIO path.
// Either way, accepted IRPs return STATUS_PENDING and complete through
// IoCompleteRequest().
// ---------------------------------------------------------------------------
static KSTATUS IdeDispatchReadWrite(DEVICE_OBJECT *device, IRP *irp) {
    (void)device;
    if (!irp || !irp->Buffer || irp->SectorCount == 0) return STATUS_INVALID_PARAMETER;
    if (irp->Lba + irp->SectorCount > gSectorCount) return STATUS_INVALID_PARAMETER;

    irp->Status = STATUS_FAILURE;
    irp->Information = 0;

    if (gDmaEnabled) {
        IrpQueueEnqueue(&gIrpQueue, irp);
        return STATUS_PENDING;
    }

    bool write = !!(irp->Flags & IRP_FLAG_WRITE);
    KSTATUS st = write ? IdeWrite(NULL, irp->Lba, irp->SectorCount, irp->Buffer)
                       : IdeRead (NULL, irp->Lba, irp->SectorCount, irp->Buffer);
    IoCompleteRequest(irp, st, KSUCCESS(st) ? ((size_t)irp->SectorCount * 512u) : 0);
    return STATUS_PENDING;
}

// ---------------------------------------------------------------------------
// IdeDmaInit — locate the PCI IDE controller, grab BAR4, allocate the PRDT,
// create sync objects, and register IRQ 14.
//
// Returns true on success. On any failure, logs a message and returns false;
// dispatch will complete IRPs inline through PIO.
// ---------------------------------------------------------------------------

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

    // ---- IDENTIFY drive (same as before) ----------------------------------
    __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) {
        DbgPrint("IDE: No drive on primary channel\n");
        return STATUS_NOT_FOUND;
    }

    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: Not a standard ATA drive (ATAPI?)\n");
        return STATUS_NOT_SUPPORTED;
    }

    timeout = 10000;
    while (true) {
        uint8_t st = __inbyte(IDE_STATUS);
        if (st & IDE_ERR_F) return STATUS_DEVICE_ERROR;
        if (st & 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);

    gSupportsLba48 = (ident[83] & (1 << 10)) != 0;
    if (gSupportsLba48) {
        gSectorCount = *(uint64_t *)&ident[100];
        DbgPrint("IDE: LBA48, %llu sectors\n", gSectorCount);
    } else {
        gSectorCount = ((uint32_t)ident[61] << 16) | ident[60];
        DbgPrint("IDE: LBA28, %u sectors\n", (uint32_t)gSectorCount);
    }
    if (gSectorCount == 0) return STATUS_NOT_FOUND;

    // ---- Register definitive DEVICE_OBJECT --------------------------------
    memset(&gDriver, 0, sizeof(gDriver));
    gDriver.Name = "ide";
    gDriver.MajorFunction[IRP_MJ_READ] = IdeDispatchReadWrite;
    gDriver.MajorFunction[IRP_MJ_WRITE] = IdeDispatchReadWrite;

    KSTATUS st = IoCreateDevice(&gDriver, 0, &gDevice);
    if (!KSUCCESS(st)) return st;
    gDevice->SectorCount = gSectorCount;
    gDevice->AlignmentRequirement = 512;

    // ---- Try to initialise Bus Master DMA ---------------------------------
    gDmaEnabled = IdeDmaInit();

    DbgPrint("IDE: Registered (%s)\n", gDmaEnabled ? "DMA" : "PIO fallback");
    return IoRegisterDiskDevice(0, gDevice);
}