#include "ata.h"
#include "io.h"      // your existing inb/outb/insw
#include "stdio.h"
#include "memory.h"
#include <stdbool.h>





#define ATA_PRIMARY_DATA       0x1F0
#define ATA_PRIMARY_ERROR      0x1F1
#define ATA_PRIMARY_SECCOUNT   0x1F2
#define ATA_PRIMARY_LBALOW     0x1F3
#define ATA_PRIMARY_LBAMID     0x1F4
#define ATA_PRIMARY_LBAHIGH    0x1F5
#define ATA_PRIMARY_DRIVE      0x1F6
#define ATA_PRIMARY_STATUS     0x1F7
#define ATA_PRIMARY_COMMAND    0x1F7

#define ATA_STATUS_BSY  0x80
#define ATA_STATUS_DRDY 0x40
#define ATA_STATUS_DRQ  0x08
#define ATA_STATUS_ERR  0x01

#define ATA_CMD_IDENTIFY      0xEC
#define ATA_CMD_READ_SECTORS  0x20

static void ata_wait_ready(void)
{
    while (x86_64_inb(ATA_PRIMARY_STATUS) & ATA_STATUS_BSY) {}
    while (!(x86_64_inb(ATA_PRIMARY_STATUS) & ATA_STATUS_DRDY)) {}
}

static bool ata_poll(void)
{
    // 400 ns delay
    for (int i = 0; i < 4; i++)
        x86_64_inb(ATA_PRIMARY_STATUS);

    uint8_t status;

    while ((status = x86_64_inb(ATA_PRIMARY_STATUS)) & ATA_STATUS_BSY);
    if (status & ATA_STATUS_ERR) return false;

    while (!(status & ATA_STATUS_DRQ)) {
        status = x86_64_inb(ATA_PRIMARY_STATUS);
        if (status & ATA_STATUS_ERR) return false;
    }

    return true;
}

void ata_init(void)
{
    x86_64_outb(0x3F6, 0x02);          // disable interrupts on primary controller
    ata_wait_ready();
    x86_64_outb(ATA_PRIMARY_DRIVE, 0xA0);   // select master
    ata_wait_ready();
    printf("ATA: Primary master (hda) initialized\n");
}

void ata_identify(void)
{
    static uint16_t buffer[256];

    x86_64_outb(ATA_PRIMARY_DRIVE, 0xA0);
    x86_64_outb(ATA_PRIMARY_SECCOUNT, 0);
    x86_64_outb(ATA_PRIMARY_LBALOW, 0);
    x86_64_outb(ATA_PRIMARY_LBAMID, 0);
    x86_64_outb(ATA_PRIMARY_LBAHIGH, 0);
    x86_64_outb(ATA_PRIMARY_COMMAND, ATA_CMD_IDENTIFY);

    if (x86_64_inb(ATA_PRIMARY_STATUS) == 0) {
        printf("ATA: No primary master drive detected\n");
        return;
    }

    if (!ata_poll()) {
        printf("ATA poll failed at ATA_IDENTIFY\n");
    }

    if (x86_64_inb(ATA_PRIMARY_STATUS) & ATA_STATUS_ERR) {
        printf("ATA: Identify error\n");
        return;
    }

    x86_64_insw(ATA_PRIMARY_DATA, buffer, 256);

    printf("ATA: Disk model: ");
    for (int i = 27; i <= 46; i++) {
        uint16_t w = buffer[i];
        char high = w >> 8;
        char low  = w & 0xFF;
        printf("%c%c", high, low); // swap bytes
    }
    printf("\n");
    printf("ATA: Total sectors: %lu\n", (uint64_t)buffer[60] | ((uint64_t)buffer[61] << 16));
}

bool ata_read_sectors(uint64_t lba, uint8_t sector_count, void* buffer)
{
    if (sector_count == 0) return true;

    ata_wait_ready();

    x86_64_outb(ATA_PRIMARY_DRIVE, 0xE0 | ((lba >> 24) & 0x0F));
    x86_64_outb(ATA_PRIMARY_SECCOUNT, sector_count);
    x86_64_outb(ATA_PRIMARY_LBALOW,  (uint8_t)lba);
    x86_64_outb(ATA_PRIMARY_LBAMID,  (uint8_t)(lba >> 8));
    x86_64_outb(ATA_PRIMARY_LBAHIGH, (uint8_t)(lba >> 16));

    x86_64_outb(ATA_PRIMARY_COMMAND, ATA_CMD_READ_SECTORS);

    uint16_t* buf = buffer;
    for (uint8_t i = 0; i < sector_count; i++) {
        if (!ata_poll()) {
            printf("ATA poll failed at ATA_read_sectors\n");
        }

        if (x86_64_inb(ATA_PRIMARY_STATUS) & ATA_STATUS_ERR) {
            printf("ATA read error at LBA %lu\n", lba + i);
            return false;
        }

        x86_64_insw(ATA_PRIMARY_DATA, buf, 256);   // 512 bytes
        buf += 256;
    }

    return true;
}

bool ata_write_sectors(uint64_t lba, uint8_t sector_count, const void* buffer) {
    if (sector_count == 0) return true;



    ata_wait_ready();
    x86_64_outb(ATA_PRIMARY_DRIVE, 0xE0 | ((lba >> 24) & 0x0F));
    x86_64_outb(ATA_PRIMARY_SECCOUNT, sector_count);
    x86_64_outb(ATA_PRIMARY_LBALOW, (uint8_t)lba);
    x86_64_outb(ATA_PRIMARY_LBAMID, (uint8_t)(lba >> 8));
    x86_64_outb(ATA_PRIMARY_LBAHIGH, (uint8_t)(lba >> 16));
    x86_64_outb(ATA_PRIMARY_COMMAND, 0x30); // write sectors

    const uint16_t* buf = buffer;
    for (uint8_t i = 0; i < sector_count; i++) {
        if (!ata_poll()) {

            return false;
        }
        x86_64_outsw(ATA_PRIMARY_DATA, buf, 256);
        buf += 256;
    }



    return true;
}
