/**
 * @file floppy.c
 * @brief Floppy disk controller driver implementation.
 *
 * Drives the Intel 82077AA-compatible floppy disk controller using
 * ISA DMA channel 2 for data transfers and IRQ 6 for completion
 * notification.
 *
 * Supports 1.44 MB 3.5" HD floppies. The driver detects drives via
 * CMOS register 0x10, initializes the controller, and registers each
 * detected drive with the devicefs subsystem as a "floppy" class
 * block device (floppy1, floppy2, ...).
 *
 * LBA addressing is converted to CHS for the floppy controller.
 */

#include "floppy.h"
#include "port_io.h"
#include "pic.h"
#include "devicefs.h"
#include "driver.h"
#include <string.h>

/* Debug print helpers defined in kernel.c */
extern void offset_print(const char *str);
extern void print_hex(uint32_t val);

/* ================================================================
 * Global state
 * ================================================================ */

/** Detected floppy drives. */
static floppy_drive_t floppy_drives[FLOPPY_MAX_DRIVES];

/** Volatile flag set by IRQ 6 handler. */
static volatile int floppy_irq_received = 0;

/** DMA bounce buffer — must be in the first 16 MB of physical memory
 *  and must not cross a 64 KB boundary. We use a static buffer that
 *  the linker places in the kernel BSS (which lives in low memory). */
static uint8_t dma_buffer[FLOPPY_SECTOR_SIZE] __attribute__((aligned(512)));

/* ================================================================
 * Low-level helpers
 * ================================================================ */

/**
 * Read CMOS register.
 *
 * @param reg CMOS register address (0x00-0x7F).
 * @return Register value.
 */
static uint8_t cmos_read(uint8_t reg) {
    outb(0x70, reg);
    /* Small delay for CMOS to respond */
    inb(0x80);
    return inb(0x71);
}

/**
 * Wait for IRQ 6 with timeout.
 * The floppy controller fires IRQ 6 on command completion.
 *
 * If interrupts are currently disabled (e.g. during early init before
 * the kernel calls STI), we temporarily enable them so IRQ 6 can
 * actually fire, then restore the previous interrupt state.
 *
 * @return 0 on success (IRQ received), -1 on timeout.
 */
static int fdc_wait_irq(void) {
    /* Check whether interrupts are currently enabled (IF flag) */
    uint32_t eflags;
    asm volatile("pushfl; popl %0" : "=r"(eflags));
    int need_sti = !(eflags & 0x200);

    if (need_sti) {
        asm volatile("sti");
    }

    int timeout = 5000000;
    while (!floppy_irq_received && --timeout > 0) {
        /* Busy wait — in a real OS we'd use a proper sleep/wait */
        asm volatile("pause");
    }

    if (need_sti) {
        asm volatile("cli");
    }

    if (timeout == 0) {
        offset_print("  FLOPPY: IRQ timeout\n");
        return -1;
    }
    floppy_irq_received = 0;
    return 0;
}

/**
 * Write a byte to the FDC FIFO register.
 * Waits for the controller to be ready (MSR RQM=1, DIO=0).
 *
 * @param val Byte to write.
 * @return 0 on success, -1 on timeout.
 */
static int fdc_write_byte(uint8_t val) {
    int timeout = 500000;
    while (--timeout > 0) {
        uint8_t msr = inb(FDC_MSR);
        if ((msr & (MSR_RQM | MSR_DIO)) == MSR_RQM) {
            outb(FDC_FIFO, val);
            return 0;
        }
    }
    offset_print("  FLOPPY: write timeout\n");
    return -1;
}

/**
 * Read a byte from the FDC FIFO register.
 * Waits for the controller to have data ready (MSR RQM=1, DIO=1).
 *
 * @return Byte read, or -1 on timeout.
 */
static int fdc_read_byte(void) {
    int timeout = 500000;
    while (--timeout > 0) {
        uint8_t msr = inb(FDC_MSR);
        if ((msr & (MSR_RQM | MSR_DIO)) == (MSR_RQM | MSR_DIO)) {
            return (int)inb(FDC_FIFO);
        }
    }
    offset_print("  FLOPPY: read timeout\n");
    return -1;
}

/**
 * Issue a SENSE INTERRUPT command and read the result.
 *
 * @param st0_out Output: ST0 byte.
 * @param cyl_out Output: current cylinder.
 */
static void fdc_sense_interrupt(uint8_t *st0_out, uint8_t *cyl_out) {
    fdc_write_byte(FDC_CMD_SENSE_INT);
    int st0 = fdc_read_byte();
    int cyl = fdc_read_byte();
    if (st0_out) *st0_out = (uint8_t)(st0 >= 0 ? st0 : 0xFF);
    if (cyl_out) *cyl_out = (uint8_t)(cyl >= 0 ? cyl : 0xFF);
}

/* ================================================================
 * Motor control
 * ================================================================ */

/**
 * Turn the motor on for a drive.
 *
 * @param drive Drive number (0 or 1).
 */
static void fdc_motor_on(uint8_t drive) {
    if (drive >= FLOPPY_MAX_DRIVES) return;
    if (floppy_drives[drive].motor_on) return;

    uint8_t dor = drive | DOR_RESET | DOR_DMA_IRQ |
                  (DOR_MOTOR_A << drive);
    outb(FDC_DOR, dor);

    /* Wait ~500 ms for motor spin-up (busy wait in timer ticks) */
    for (volatile int i = 0; i < 2000000; i++) {
        asm volatile("pause");
    }
    floppy_drives[drive].motor_on = 1;
}

/**
 * Turn the motor off for a drive.
 *
 * @param drive Drive number (0 or 1).
 */
static void fdc_motor_off(uint8_t drive) {
    if (drive >= FLOPPY_MAX_DRIVES) return;
    uint8_t dor = drive | DOR_RESET | DOR_DMA_IRQ;
    outb(FDC_DOR, dor);
    floppy_drives[drive].motor_on = 0;
}

/* ================================================================
 * DMA setup for ISA DMA channel 2
 * ================================================================ */

/**
 * Set up ISA DMA channel 2 for a floppy transfer.
 *
 * The floppy controller uses DMA channel 2. We configure the DMA
 * controller with the physical address of our bounce buffer and the
 * transfer length.
 *
 * @param addr Physical address of the DMA buffer (must be < 16 MB).
 * @param len  Transfer length in bytes (minus 1 is written to count).
 * @param read 1 for read (disk→memory), 0 for write (memory→disk).
 */
static void dma_setup(uint32_t addr, uint16_t len, int read) {
    /* Mask DMA channel 2 */
    outb(0x0A, 0x06);

    /* Reset flip-flop */
    outb(0x0C, 0xFF);

    /* Mode: single transfer, auto-init disabled, increment, channel 2
     * Read = 0x46 (single, addr incr, write-to-memory = FDC read)
     * Write = 0x4A (single, addr incr, read-from-memory = FDC write) */
    outb(0x0B, read ? 0x46 : 0x4A);

    /* Address (low 16 bits via channel 2 address reg 0x04) */
    outb(0x04, (uint8_t)(addr & 0xFF));
    outb(0x04, (uint8_t)((addr >> 8) & 0xFF));

    /* Page register for channel 2 (port 0x81) — bits 16-23 of address */
    outb(0x81, (uint8_t)((addr >> 16) & 0xFF));

    /* Count (len - 1, low byte first) */
    outb(0x05, (uint8_t)((len - 1) & 0xFF));
    outb(0x05, (uint8_t)(((len - 1) >> 8) & 0xFF));

    /* Unmask DMA channel 2 */
    outb(0x0A, 0x02);
}

/* ================================================================
 * CHS conversion
 * ================================================================ */

/**
 * Convert an LBA address to CHS for a 1.44 MB floppy.
 *
 * @param lba     Logical block address.
 * @param cyl_out Output: cylinder number.
 * @param head_out Output: head number.
 * @param sec_out  Output: sector number (1-based).
 */
static void lba_to_chs(uint32_t lba, uint8_t *cyl_out,
                        uint8_t *head_out, uint8_t *sec_out) {
    *cyl_out  = (uint8_t)(lba / (FLOPPY_HEADS * FLOPPY_SECTORS_PER_TRACK));
    uint32_t tmp = lba % (FLOPPY_HEADS * FLOPPY_SECTORS_PER_TRACK);
    *head_out = (uint8_t)(tmp / FLOPPY_SECTORS_PER_TRACK);
    *sec_out  = (uint8_t)((tmp % FLOPPY_SECTORS_PER_TRACK) + 1);
}

/* ================================================================
 * Seek and calibrate
 * ================================================================ */

/**
 * Recalibrate a drive (seek to track 0).
 *
 * @param drive Drive number.
 * @return 0 on success, -1 on failure.
 */
static int fdc_recalibrate(uint8_t drive) {
    floppy_irq_received = 0;
    fdc_write_byte(FDC_CMD_RECALIBRATE);
    fdc_write_byte(drive);

    if (fdc_wait_irq() != 0) return -1;

    uint8_t st0, cyl;
    fdc_sense_interrupt(&st0, &cyl);

    if (cyl != 0) {
        offset_print("  FLOPPY: recalibrate failed, cyl=");
        print_hex(cyl);
        offset_print("\n");
        return -1;
    }

    floppy_drives[drive].current_track = 0;
    return 0;
}

/**
 * Seek to a specific cylinder.
 *
 * @param drive Drive number.
 * @param cyl   Target cylinder.
 * @return 0 on success, -1 on failure.
 */
static int fdc_seek(uint8_t drive, uint8_t cyl) {
    if (floppy_drives[drive].current_track == cyl) return 0;

    floppy_irq_received = 0;
    fdc_write_byte(FDC_CMD_SEEK);
    fdc_write_byte((0 << 2) | drive);  /* Head 0, drive N */
    fdc_write_byte(cyl);

    if (fdc_wait_irq() != 0) return -1;

    uint8_t st0, res_cyl;
    fdc_sense_interrupt(&st0, &res_cyl);

    if (res_cyl != cyl) {
        offset_print("  FLOPPY: seek failed, wanted cyl=");
        print_hex(cyl);
        offset_print(" got=");
        print_hex(res_cyl);
        offset_print("\n");
        return -1;
    }

    floppy_drives[drive].current_track = cyl;
    return 0;
}

/* ================================================================
 * Read / Write sector
 * ================================================================ */

/**
 * Read a single sector from a floppy drive using DMA.
 *
 * @param drive Drive number (0 or 1).
 * @param lba   Logical sector address.
 * @param buf   Output buffer (512 bytes).
 * @return 0 on success, -1 on failure.
 */
static int fdc_read_sector(uint8_t drive, uint32_t lba, void *buf) {
    if (lba >= FLOPPY_TOTAL_SECTORS) return -1;

    uint8_t cyl, head, sec;
    lba_to_chs(lba, &cyl, &head, &sec);

    fdc_motor_on(drive);

    /* Seek to the correct cylinder */
    if (fdc_seek(drive, cyl) != 0) {
        fdc_motor_off(drive);
        return -1;
    }

    /* Set up DMA for read (disk → memory) */
    uint32_t dma_addr = (uint32_t)dma_buffer;
    dma_setup(dma_addr, FLOPPY_SECTOR_SIZE, 1);

    /* Issue READ DATA command */
    floppy_irq_received = 0;
    fdc_write_byte(FDC_CMD_READ_DATA);
    fdc_write_byte((head << 2) | drive);
    fdc_write_byte(cyl);
    fdc_write_byte(head);
    fdc_write_byte(sec);
    fdc_write_byte(2);    /* Sector size code: 2 = 512 bytes */
    fdc_write_byte(FLOPPY_SECTORS_PER_TRACK); /* End of track */
    fdc_write_byte(0x1B); /* GPL: gap length */
    fdc_write_byte(0xFF); /* DTL: unused when sector size != 0 */

    /* Wait for command completion */
    if (fdc_wait_irq() != 0) {
        fdc_motor_off(drive);
        return -1;
    }

    /* Read result phase (7 bytes) */
    uint8_t st0 = (uint8_t)fdc_read_byte();
    uint8_t st1 = (uint8_t)fdc_read_byte();
    uint8_t st2 = (uint8_t)fdc_read_byte();
    (void)fdc_read_byte(); /* C */
    (void)fdc_read_byte(); /* H */
    (void)fdc_read_byte(); /* R */
    (void)fdc_read_byte(); /* N */

    /* Check for errors */
    if ((st0 & 0xC0) != 0 || st1 != 0 || st2 != 0) {
        offset_print("  FLOPPY: read error st0=");
        print_hex(st0);
        offset_print(" st1=");
        print_hex(st1);
        offset_print(" st2=");
        print_hex(st2);
        offset_print("\n");
        fdc_motor_off(drive);
        return -1;
    }

    /* Copy from DMA buffer to caller's buffer */
    memcpy(buf, dma_buffer, FLOPPY_SECTOR_SIZE);

    fdc_motor_off(drive);
    return 0;
}

/**
 * Write a single sector to a floppy drive using DMA.
 *
 * @param drive Drive number (0 or 1).
 * @param lba   Logical sector address.
 * @param buf   Input buffer (512 bytes).
 * @return 0 on success, -1 on failure.
 */
static int fdc_write_sector(uint8_t drive, uint32_t lba, const void *buf) {
    if (lba >= FLOPPY_TOTAL_SECTORS) return -1;

    uint8_t cyl, head, sec;
    lba_to_chs(lba, &cyl, &head, &sec);

    fdc_motor_on(drive);

    if (fdc_seek(drive, cyl) != 0) {
        fdc_motor_off(drive);
        return -1;
    }

    /* Copy data to DMA buffer */
    memcpy(dma_buffer, buf, FLOPPY_SECTOR_SIZE);

    /* Set up DMA for write (memory → disk) */
    uint32_t dma_addr = (uint32_t)dma_buffer;
    dma_setup(dma_addr, FLOPPY_SECTOR_SIZE, 0);

    /* Issue WRITE DATA command */
    floppy_irq_received = 0;
    fdc_write_byte(FDC_CMD_WRITE_DATA);
    fdc_write_byte((head << 2) | drive);
    fdc_write_byte(cyl);
    fdc_write_byte(head);
    fdc_write_byte(sec);
    fdc_write_byte(2);    /* Sector size code: 512 bytes */
    fdc_write_byte(FLOPPY_SECTORS_PER_TRACK);
    fdc_write_byte(0x1B);
    fdc_write_byte(0xFF);

    if (fdc_wait_irq() != 0) {
        fdc_motor_off(drive);
        return -1;
    }

    /* Read result phase */
    uint8_t st0 = (uint8_t)fdc_read_byte();
    uint8_t st1 = (uint8_t)fdc_read_byte();
    uint8_t st2 = (uint8_t)fdc_read_byte();
    (void)fdc_read_byte();
    (void)fdc_read_byte();
    (void)fdc_read_byte();
    (void)fdc_read_byte();

    if ((st0 & 0xC0) != 0 || st1 != 0 || st2 != 0) {
        offset_print("  FLOPPY: write error st0=");
        print_hex(st0);
        offset_print(" st1=");
        print_hex(st1);
        offset_print(" st2=");
        print_hex(st2);
        offset_print("\n");
        fdc_motor_off(drive);
        return -1;
    }

    fdc_motor_off(drive);
    return 0;
}

/* ================================================================
 * Devicefs block operations
 * ================================================================ */

/**
 * Read sectors from a floppy drive.
 * dev_data is a pointer to the floppy_drive_t.
 */
static int floppy_block_read(void *dev_data, uint32_t lba,
                              uint32_t count, void *buf) {
    floppy_drive_t *drv = (floppy_drive_t *)dev_data;
    if (!drv || !drv->present) return -1;

    uint8_t *p = (uint8_t *)buf;
    for (uint32_t i = 0; i < count; i++) {
        if (fdc_read_sector(drv->drive_num, lba + i, p) != 0) {
            return -1;
        }
        p += FLOPPY_SECTOR_SIZE;
    }
    return 0;
}

/**
 * Write sectors to a floppy drive.
 */
static int floppy_block_write(void *dev_data, uint32_t lba,
                               uint32_t count, const void *buf) {
    floppy_drive_t *drv = (floppy_drive_t *)dev_data;
    if (!drv || !drv->present) return -1;

    const uint8_t *p = (const uint8_t *)buf;
    for (uint32_t i = 0; i < count; i++) {
        if (fdc_write_sector(drv->drive_num, lba + i, p) != 0) {
            return -1;
        }
        p += FLOPPY_SECTOR_SIZE;
    }
    return 0;
}

/**
 * Get floppy sector size (always 512).
 */
static uint32_t floppy_block_sector_size(void *dev_data) {
    (void)dev_data;
    return FLOPPY_SECTOR_SIZE;
}

/**
 * Get total sector count for a 1.44 MB floppy.
 */
static uint32_t floppy_block_sector_count(void *dev_data) {
    (void)dev_data;
    return FLOPPY_TOTAL_SECTORS;
}

/** Block device operations for floppy drives. */
static devicefs_block_ops_t floppy_block_ops = {
    .read_sectors  = floppy_block_read,
    .write_sectors = floppy_block_write,
    .sector_size   = floppy_block_sector_size,
    .sector_count  = floppy_block_sector_count,
};

/* ================================================================
 * Controller initialization
 * ================================================================ */

/**
 * Reset the floppy disk controller.
 *
 * Toggles the RESET bit in the DOR and re-initializes the controller.
 *
 * @return 0 on success, -1 on failure.
 */
static int fdc_reset(void) {
    /* Enter reset state */
    outb(FDC_DOR, 0x00);

    /* Small delay */
    for (volatile int i = 0; i < 100000; i++) {
        asm volatile("pause");
    }

    /* Exit reset: enable DMA/IRQ, select drive 0, motor off */
    floppy_irq_received = 0;
    outb(FDC_DOR, DOR_RESET | DOR_DMA_IRQ);

    /* Controller generates IRQ 6 after reset */
    if (fdc_wait_irq() != 0) {
        offset_print("  FLOPPY: no IRQ after reset\n");
        return -1;
    }

    /* Sense interrupt for each of the 4 possible drives */
    for (int i = 0; i < 4; i++) {
        uint8_t st0, cyl;
        fdc_sense_interrupt(&st0, &cyl);
    }

    /* Set data rate to 500 kbps (for 1.44 MB HD floppies) */
    outb(FDC_CCR, 0x00);

    /* SPECIFY command: step rate = 3 ms, head unload = 240 ms,
     * head load = 16 ms, non-DMA mode = 0 */
    fdc_write_byte(FDC_CMD_SPECIFY);
    fdc_write_byte(0xDF); /* SRT=3ms, HUT=240ms */
    fdc_write_byte(0x02); /* HLT=16ms, NDMA=0 */

    return 0;
}

/* ================================================================
 * Driver framework integration
 * ================================================================ */

/**
 * Probe for floppy controller.
 * Reads CMOS to determine if any floppy drives are configured.
 */
static driver_probe_result_t floppy_probe(void) {
    uint8_t cmos_floppy = cmos_read(0x10);
    uint8_t drive_a = (cmos_floppy >> 4) & 0x0F;
    uint8_t drive_b = cmos_floppy & 0x0F;

    if (drive_a == CMOS_FLOPPY_NONE && drive_b == CMOS_FLOPPY_NONE) {
        offset_print("  FLOPPY: no drives in CMOS\n");
        return DRIVER_PROBE_NOT_FOUND;
    }

    return DRIVER_PROBE_OK;
}

/**
 * Initialize the floppy controller and register detected drives.
 */
static int floppy_driver_init(void) {
    memset(floppy_drives, 0, sizeof(floppy_drives));

    /* Read CMOS for drive types */
    uint8_t cmos_floppy = cmos_read(0x10);
    uint8_t types[2] = {
        (cmos_floppy >> 4) & 0x0F,
        cmos_floppy & 0x0F,
    };

    /* Unmask IRQ 6 in the PIC */
    pic_clear_mask(6);

    /* Reset the controller */
    if (fdc_reset() != 0) {
        offset_print("  FLOPPY: controller reset failed\n");
        return -1;
    }

    /* Register each detected drive */
    int found = 0;
    for (int i = 0; i < 2; i++) {
        if (types[i] == CMOS_FLOPPY_NONE) continue;

        /* We only fully support 1.44 MB for now, but register others too */
        floppy_drives[i].drive_num = (uint8_t)i;
        floppy_drives[i].cmos_type = types[i];
        floppy_drives[i].present = 1;
        floppy_drives[i].motor_on = 0;
        floppy_drives[i].current_track = 0xFF; /* Unknown */

        const char *type_str;
        switch (types[i]) {
        case CMOS_FLOPPY_360K:  type_str = "360K 5.25\""; break;
        case CMOS_FLOPPY_1200K: type_str = "1.2M 5.25\""; break;
        case CMOS_FLOPPY_720K:  type_str = "720K 3.5\""; break;
        case CMOS_FLOPPY_1440K: type_str = "1.44M 3.5\""; break;
        case CMOS_FLOPPY_2880K: type_str = "2.88M 3.5\""; break;
        default:                type_str = "unknown"; break;
        }

        offset_print("  FLOPPY: drive ");
        print_hex(i);
        offset_print(": ");
        offset_print(type_str);
        offset_print("\n");

        /* Recalibrate the drive */
        fdc_motor_on((uint8_t)i);
        fdc_recalibrate((uint8_t)i);
        fdc_motor_off((uint8_t)i);

        /* Register with devicefs */
        devicefs_register_block("floppy", &floppy_block_ops,
                                &floppy_drives[i]);
        found++;
    }

    if (found == 0) {
        offset_print("  FLOPPY: no usable drives found\n");
    } else {
        offset_print("  FLOPPY: ");
        print_hex((uint32_t)found);
        offset_print(" drive(s) initialized\n");
    }

    return 0;
}

/* ================================================================
 * IRQ handler
 * ================================================================ */

void floppy_irq(void) {
    floppy_irq_received = 1;
}

/* ================================================================
 * Public API & driver registration
 * ================================================================ */

int floppy_init(void) {
    return floppy_driver_init();
}

/** Driver descriptor. */
static const driver_t floppy_driver = {
    .name  = "floppy",
    .probe = floppy_probe,
    .init  = floppy_driver_init,
};

REGISTER_DRIVER(floppy_driver);