// ide.c — IDE/ATA driver: PIO (sync fallback) + Bus Master DMA (async)
//
// Architecture
// ------------
// The DMA path mirrors AHCI exactly:
//   SubmitIrp   → enqueues IRP + signals gIrpSemaphore → returns STATUS_PENDING
//   IdeWorker   → dequeues IRP, calls IdeDmaRw(), fires completion hooks
//   IdeIsr      → clears BM status, signals gDmaEvent to unblock the worker
//
// Graceful degradation: if Bus Master init fails (no PCI IDE controller,
// BAR4 not assigned, PRDT allocation above 4 GB, etc.) SubmitIrp is left
// NULL. HalDiskDispatch sees NULL and falls back to the synchronous PIO
// path with no changes at the call sites.
//
// PRDT layout
// -----------
// One page allocated via MmAllocateContiguous. We build up to PRDT_MAX_ENTRIES
// entries by walking the transfer buffer in page-sized chunks and looking up
// each page's physical address via MmGetPhysicalAddress. Adjacent physical
// pages are coalesced into a single PRD entry where possible.
//
// LBA48 vs LBA28
// --------------
// DMA commands exist in both modes: CMD_READ/WRITE_DMA_EXT (0x25/0x35) for
// LBA48 and CMD_READ/WRITE_DMA (0xC8/0xCA) for LBA28. The existing
// gSupportsLba48 flag picks the right variant.

#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>
#include <ke_irp.h>
#include <ke_event.h>
#include <ke_semaphore.h>
#include <ke_thread.h>
#include <ke_irq.h>
#include <mm_phys.h>
#include <mm_virt.h>

// ---------------------------------------------------------------------------
// IDE register map (primary channel, base 0x1F0)
// ---------------------------------------------------------------------------
#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

// ATA commands
#define CMD_READ_EXT       0x24   // PIO  48-bit read
#define CMD_WRITE_EXT      0x34   // PIO  48-bit write
#define CMD_READ_DMA_EXT   0x25   // DMA  48-bit read
#define CMD_WRITE_DMA_EXT  0x35   // DMA  48-bit write
#define CMD_READ_DMA       0xC8   // DMA  28-bit read
#define CMD_WRITE_DMA      0xCA   // DMA  28-bit write
#define CMD_IDENT          0xEC

// IDE primary IRQ
#define IDE_IRQ_PRIMARY    14

// ---------------------------------------------------------------------------
// Bus Master IDE register offsets (relative to gBmBase, primary channel)
// ---------------------------------------------------------------------------
#define BM_CMD_REG     0x00   // command
#define BM_STATUS_REG  0x02   // status
#define BM_PRDT_REG    0x04   // PRDT physical address (32-bit)

// BM_CMD_REG bits
#define BM_CMD_START   0x01   // start DMA transfer
#define BM_CMD_READ    0x08   // 1 = disk→memory (read), 0 = memory→disk (write)

// BM_STATUS_REG bits
#define BM_STATUS_ACTIVE  0x01   // DMA active
#define BM_STATUS_ERROR   0x02   // DMA error (write 1 to clear)
#define BM_STATUS_IRQ     0x04   // interrupt pending (write 1 to clear)

// ---------------------------------------------------------------------------
// Physical Region Descriptor (PRDT entry, 8 bytes)
// The PRDT itself must not cross a 64 KB boundary and must be 4-byte aligned.
// We allocate a full page so both constraints are trivially satisfied.
// ---------------------------------------------------------------------------
#pragma pack(push, 1)
typedef struct {
    uint32_t physAddr;   // physical base of this memory region
    uint16_t byteCount;  // byte count; 0 means 64 KB
    uint16_t flags;      // bit 15: End Of Table
} PRD;
#pragma pack(pop)

#define PRD_EOT           0x8000u
// 512 entries × 8 bytes = 4096 bytes = exactly one page; supports 512×64 KB = 32 MB
#define PRDT_MAX_ENTRIES  512

// ---------------------------------------------------------------------------
// Module state
// ---------------------------------------------------------------------------
static const HAL_API *gHal;
static DISK_DRIVER    gDrv;
static uint64_t       gSectorCount   = 0;
static bool           gSupportsLba48 = false;

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

// Async worker synchronisation
static IRP_QUEUE  gIrpQueue;
static SEMAPHORE *gIrpSemaphore = NULL;
static EVENT     *gDmaEvent     = NULL;  // auto-reset; ISR signals after each xfer

// ---------------------------------------------------------------------------
// 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%llx has phys 0x%llx > 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 ISR (IRQ 14 — primary channel)
//
// Must be fast — runs with interrupts disabled. We only clear the hardware
// interrupt bits here and signal gDmaEvent to wake the worker thread.
// The worker reads the final status and fires IRP completion hooks.
// ---------------------------------------------------------------------------
static 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 the worker thread; it will read the final BM status and complete
    // the IRP. gDmaEvent is auto-reset so exactly one waiter is released.
    KeSignalEvent(gDmaEvent);
}

// ---------------------------------------------------------------------------
// 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);

    // ---- 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.
// ---------------------------------------------------------------------------
static void IdeWorker(uintptr_t arg) {
    (void)arg;
    while (true) {
        KeWaitSemaphore(gIrpSemaphore);

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

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

        if (irp->CompletionEvent) KeSignalEvent(irp->CompletionEvent);
        if (irp->CompletionFn)   irp->CompletionFn(irp, irp->CompletionCtx);
    }
}

// ---------------------------------------------------------------------------
// IdeSubmitIrp — async entry point wired into DISK_DRIVER.SubmitIrp.
// Enqueues the IRP and immediately returns STATUS_PENDING.
// ---------------------------------------------------------------------------
static KSTATUS IdeSubmitIrp(void *ctx, IRP *irp) {
    (void)ctx;
    IrpQueueEnqueue(&gIrpQueue, irp);
    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;
// DriverEntry then leaves DISK_DRIVER.SubmitIrp = NULL so the dispatch layer
// falls back to the synchronous PIO path automatically.
// ---------------------------------------------------------------------------
static bool IdeDmaInit(void) {
    // ---- Find PCI IDE controller (class 0x01, subclass 0x01) --------------
    uint8_t bus, dev, fn;
    if (!gHal->Pci.FindClass(0x01, 0x01, &bus, &dev, &fn)) {
        DbgPrint("IDE DMA: No PCI IDE controller found — using PIO\n");
        return false;
    }

    // ---- Read BAR4 (Bus Master base, I/O space) ---------------------------
    // PCI config offset 0x20. Bit 0 must be 1 for I/O BAR.
    uint32_t bar4 = gHal->Pci.Read32(bus, dev, fn, 0x20);
    if (!(bar4 & 0x1)) {
        DbgPrint("IDE DMA: BAR4 is a memory BAR (0x%08x) — unexpected, using PIO\n", bar4);
        return false;
    }
    uint16_t bmBase = (uint16_t)(bar4 & 0xFFFC);
    if (bmBase == 0) {
        DbgPrint("IDE DMA: BAR4 is zero (firmware didn't assign it) — using PIO\n");
        return false;
    }
    DbgPrint("IDE DMA: Bus Master base = 0x%04X (bus %u, dev %u, fn %u)\n",
             bmBase, bus, dev, fn);

    // ---- Enable Bus Mastering in PCI Command register (bit 2) ------------
    uint32_t pciCmd = gHal->Pci.Read32(bus, dev, fn, 0x04);
    gHal->Pci.Write32(bus, dev, fn, 0x04, pciCmd | (1u << 2));

    // ---- Allocate PRDT page (must be 32-bit physical, 4-byte aligned) ----
    PHYSICAL_ADDRESS prdtPhys = MmAllocateContiguous(1);
    if (!prdtPhys) {
        DbgPrint("IDE DMA: Failed to allocate PRDT page\n");
        return false;
    }
    if (prdtPhys >> 32) {
        // MmAllocateContiguous returned memory above 4 GB — BM IDE can't
        // address it. This shouldn't happen on normal x86 systems with early
        // boot allocations but we check defensively.
        DbgPrint("IDE DMA: PRDT page is above 4 GB (0x%llx) — using PIO\n",
                 (uint64_t)prdtPhys);
        return false;
    }
    void *prdtVirt = gHal->PhysToVirt(prdtPhys);
    DbgPrint("IDE DMA: PRDT at phys 0x%08x, virt %p\n", (uint32_t)prdtPhys, prdtVirt);

    // ---- Create synchronisation objects ------------------------------------
    SEMAPHORE *sem = NULL;
    if (!KSUCCESS(KeSemaphoreCreate(0, 0x7FFFFFFF, &sem))) {
        DbgPrint("IDE DMA: Failed to create IRP semaphore\n");
        return false;
    }

    EVENT *dmaEv = NULL;
    if (!KSUCCESS(KeEventCreate(EventAutoReset, 0, &dmaEv))) {
        DbgPrint("IDE DMA: Failed to create DMA completion event\n");
        return false;
    }

    // ---- Commit module state -----------------------------------------------
    gBmBase    = bmBase;
    gPrdtPhys  = (uint32_t)prdtPhys;
    gPrdtVirt  = (PRD *)prdtVirt;
    gIrpSemaphore = sem;
    gDmaEvent     = dmaEv;

    IrpQueueInit(&gIrpQueue, sem);

    // ---- Register IRQ 14 --------------------------------------------------
    KeIrqRegister(IDE_IRQ_PRIMARY, IdeIsr);
    DbgPrint("IDE DMA: IRQ %u registered\n", IDE_IRQ_PRIMARY);

    // ---- Spawn worker thread -----------------------------------------------
    KSTATUS st = KeCreateKernelThread((uintptr_t)IdeWorker, 0,
                                      PriorityHigh, (THREAD **)NULL);
    if (!KSUCCESS(st)) {
        DbgPrint("IDE DMA: Worker thread creation failed (0x%x)\n", st);
        // Not fatal for init — the PRDT/IRQ are live but async callers will
        // never get completions. Better to surface the failure clearly.
        return false;
    }

    DbgPrint("IDE DMA: Init complete\n");
    return true;
}

// ---------------------------------------------------------------------------
// 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;

    // ---- Try to initialise Bus Master DMA ---------------------------------
    // On success, wire in IdeSubmitIrp. On failure, leave it NULL so
    // HalDiskDispatch falls back to the synchronous PIO callbacks.
    bool dmaOk = IdeDmaInit();

    gDrv.Ctx          = NULL;
    gDrv.ReadSectors  = IdeRead;
    gDrv.WriteSectors = IdeWrite;
    gDrv.SubmitIrp    = dmaOk ? IdeSubmitIrp : NULL;
    gDrv.SectorCount  = gSectorCount;

    DbgPrint("IDE: Registered (%s)\n", dmaOk ? "DMA" : "PIO fallback");
    return IoRegisterDisk(0, &gDrv);
}