/* ============================================================================
 * Pactor64 — ATA PIO (Programmed I/O) Block Device Driver
 * Primary ATA controller: ports 0x1F0-0x1F7, 0x3F6
 * ============================================================================ */
#include "../include/pactor64.h"

/* === ATA I/O Ports (Primary Bus) === */
#define ATA_DATA        0x1F0   /* Data register (R/W, 16-bit) */
#define ATA_ERROR       0x1F1   /* Error register (R) */
#define ATA_FEATURES    0x1F1   /* Features register (W) */
#define ATA_SECCOUNT    0x1F2   /* Sector count register */
#define ATA_LBA_LO      0x1F3   /* LBA low byte */
#define ATA_LBA_MID     0x1F4   /* LBA mid byte */
#define ATA_LBA_HI      0x1F5   /* LBA high byte */
#define ATA_DRIVE_HEAD  0x1F6   /* Drive/Head register */
#define ATA_STATUS      0x1F7   /* Status register (R) */
#define ATA_COMMAND     0x1F7   /* Command register (W) */
#define ATA_ALT_STATUS  0x3F6   /* Alternate status (R) */
#define ATA_DEV_CTRL    0x3F6   /* Device control (W) */

/* === ATA Status Bits === */
#define ATA_SR_BSY      0x80    /* Busy */
#define ATA_SR_DRDY     0x40    /* Drive ready */
#define ATA_SR_DF       0x20    /* Drive write fault */
#define ATA_SR_DSC      0x10    /* Drive seek complete */
#define ATA_SR_DRQ      0x08    /* Data request ready */
#define ATA_SR_CORR     0x04    /* Corrected data */
#define ATA_SR_IDX      0x02    /* Index */
#define ATA_SR_ERR      0x01    /* Error */

/* === ATA Commands === */
#define ATA_CMD_IDENTIFY    0xEC
#define ATA_CMD_READ        0x20
#define ATA_CMD_WRITE       0x30

/* === Global State === */
static int ata_drive_exists = 0;
static char ata_model[41];
static char ata_serial[21];
static uint32_t ata_sectors = 0;

/* === Helper: Wait for drive to be not busy === */
static void ata_wait_bsy(void) {
    while (inb(ATA_STATUS) & ATA_SR_BSY)
        ;
}

/* === Helper: Wait for DRQ (data ready) === */
static int ata_wait_drq(void) {
    uint8_t status;
    /* Wait up to ~1 second (poll ~1M times) */
    for (int i = 0; i < 1000000; i++) {
        status = inb(ATA_STATUS);
        if (status & ATA_SR_ERR) return -1;
        if (status & ATA_SR_DF)  return -1;
        if (status & ATA_SR_DRQ) return 0;
    }
    return -1; /* timeout */
}

/* === Helper: 400ns delay via reading alternate status === */
static void ata_delay_400ns(void) {
    inb(ATA_ALT_STATUS);
    inb(ATA_ALT_STATUS);
    inb(ATA_ALT_STATUS);
    inb(ATA_ALT_STATUS);
}

/* === Initialize ATA controller and detect drive === */
void ata_init(void) {
    kprintf("[ata] Initializing primary ATA controller...\n");

    /* Select drive 0 (master) on primary bus */
    outb(ATA_DRIVE_HEAD, 0xA0);   /* LBA mode, drive 0 */
    ata_delay_400ns();

    /* Soft reset */
    outb(ATA_DEV_CTRL, 0x04);     /* Set SRST bit */
    ata_delay_400ns();
    outb(ATA_DEV_CTRL, 0x00);     /* Clear SRST */
    ata_delay_400ns();

    /* Wait for drive to be ready */
    ata_wait_bsy();

    /* Check if drive exists by reading status */
    uint8_t status = inb(ATA_STATUS);
    if (status == 0x00 || status == 0xFF) {
        kprintf("[ata] No drive detected (status=0x%x)\n", status);
        return;
    }

    kprintf("[ata] Drive detected, status=0x%x\n", status);

    /* Send IDENTIFY command */
    outb(ATA_DRIVE_HEAD, 0xA0);   /* Drive 0, LBA */
    outb(ATA_SECCOUNT, 0);
    outb(ATA_LBA_LO, 0);
    outb(ATA_LBA_MID, 0);
    outb(ATA_LBA_HI, 0);
    outb(ATA_COMMAND, ATA_CMD_IDENTIFY);

    /* Wait for response */
    ata_delay_400ns();
    status = inb(ATA_STATUS);

    if (status == 0x00) {
        kprintf("[ata] Drive does not support IDENTIFY\n");
        return;
    }

    /* Wait for BSY to clear */
    ata_wait_bsy();

    /* Check LBA mid/high for ATAPI detection */
    if (inb(ATA_LBA_MID) != 0 || inb(ATA_LBA_HI) != 0) {
        kprintf("[ata] Not an ATA drive (ATAPI?)\n");
        return;
    }

    /* Wait for DRQ */
    if (ata_wait_drq() != 0) {
        kprintf("[ata] IDENTIFY timeout\n");
        return;
    }

    /* Read 256 words (512 bytes) of IDENTIFY data */
    uint16_t identify_data[256];
    for (int i = 0; i < 256; i++) {
        identify_data[i] = inw(ATA_DATA);
    }

    /* Extract model string (words 27-46, 40 bytes) */
    for (int i = 0; i < 20; i++) {
        ata_model[i * 2]     = (char)(identify_data[27 + i] >> 8);
        ata_model[i * 2 + 1] = (char)(identify_data[27 + i] & 0xFF);
    }
    ata_model[40] = '\0';

    /* Trim trailing spaces from model */
    for (int i = 39; i >= 0 && ata_model[i] == ' '; i--)
        ata_model[i] = '\0';

    /* Extract serial string (words 10-19, 20 bytes) */
    for (int i = 0; i < 10; i++) {
        ata_serial[i * 2]     = (char)(identify_data[10 + i] >> 8);
        ata_serial[i * 2 + 1] = (char)(identify_data[10 + i] & 0xFF);
    }
    ata_serial[20] = '\0';

    /* Trim trailing spaces from serial */
    for (int i = 19; i >= 0 && ata_serial[i] == ' '; i--)
        ata_serial[i] = '\0';

    /* Capacity: words 60-61 contain total LBA28 sectors */
    ata_sectors = ((uint32_t)identify_data[61] << 16) | identify_data[60];
    uint32_t capacity_mb = (ata_sectors * 512) / (1024 * 1024);

    ata_drive_exists = 1;

    kprintf("[ata] Model:    %s\n", ata_model);
    kprintf("[ata] Serial:   %s\n", ata_serial);
    kprintf("[ata] Sectors:  %d (%d MB)\n", ata_sectors, capacity_mb);
    kprintf("[ata] ATA PIO driver ready.\n");
}

/* === Read sectors using LBA28 PIO === */
int ata_read_sectors(uint32_t lba, uint32_t count, void *buf) {
    if (!ata_drive_exists) return -1;

    uint16_t *ptr = (uint16_t *)buf;

    for (uint32_t i = 0; i < count; i++) {
        uint32_t cur_lba = lba + i;

        ata_wait_bsy();

        /* Select drive + upper 4 bits of LBA */
        outb(ATA_DRIVE_HEAD, 0xE0 | ((cur_lba >> 24) & 0x0F));
        ata_delay_400ns();

        /* Set parameters */
        outb(ATA_SECCOUNT, 1);                      /* 1 sector at a time */
        outb(ATA_LBA_LO,  (uint8_t)(cur_lba & 0xFF));
        outb(ATA_LBA_MID, (uint8_t)((cur_lba >> 8) & 0xFF));
        outb(ATA_LBA_HI,  (uint8_t)((cur_lba >> 16) & 0xFF));

        /* Send READ SECTORS command */
        outb(ATA_COMMAND, ATA_CMD_READ);

        /* Wait for data */
        if (ata_wait_drq() != 0) {
            kprintf("[ata] Read error at LBA %d\n", cur_lba);
            return -1;
        }

        /* Read 256 words (512 bytes) */
        for (int w = 0; w < 256; w++) {
            ptr[w] = inw(ATA_DATA);
        }
        ptr += 256;
    }

    return 0;
}

/* === Write sectors using LBA28 PIO === */
int ata_write_sectors(uint32_t lba, uint32_t count, const void *buf) {
    if (!ata_drive_exists) return -1;

    const uint16_t *ptr = (const uint16_t *)buf;

    for (uint32_t i = 0; i < count; i++) {
        uint32_t cur_lba = lba + i;

        ata_wait_bsy();

        /* Select drive + upper 4 bits of LBA */
        outb(ATA_DRIVE_HEAD, 0xE0 | ((cur_lba >> 24) & 0x0F));
        ata_delay_400ns();

        /* Set parameters */
        outb(ATA_SECCOUNT, 1);
        outb(ATA_LBA_LO,  (uint8_t)(cur_lba & 0xFF));
        outb(ATA_LBA_MID, (uint8_t)((cur_lba >> 8) & 0xFF));
        outb(ATA_LBA_HI,  (uint8_t)((cur_lba >> 16) & 0xFF));

        /* Send WRITE SECTORS command */
        outb(ATA_COMMAND, ATA_CMD_WRITE);

        /* Wait for DRQ */
        if (ata_wait_drq() != 0) {
            kprintf("[ata] Write error at LBA %d\n", cur_lba);
            return -1;
        }

        /* Write 256 words (512 bytes) */
        for (int w = 0; w < 256; w++) {
            outw(ATA_DATA, ptr[w]);
        }
        ptr += 256;

        /* Flush: send CACHE FLUSH command */
        ata_wait_bsy();
        outb(ATA_COMMAND, 0xE7);  /* CACHE FLUSH */
        ata_wait_bsy();
    }

    return 0;
}

/* === Check if drive exists === */
int ata_is_present(void) {
    return ata_drive_exists;
}

/* === Get drive info === */
uint32_t ata_get_sector_count(void) {
    return ata_sectors;
}

const char *ata_get_model(void) {
    return ata_model;
}

const char *ata_get_serial(void) {
    return ata_serial;
}

/* === Raw sector read (for ext2 block device callback) === */
/* Note: This wraps ata_read_sectors with LBA/sector conversion */
int ata_block_read(uint64_t lba, uint32_t count, void *buf) {
    return ata_read_sectors((uint32_t)lba, count, buf);
}