#include "alloc.h"
#include "common.h"
#include "map.h"
#include "print.h"
#include "lib/stdio.h"
#include "map.h"
#include "lib/malloc.h"

#define PCI_CONFIG_ADDR 0xCF8
#define PCI_CONFIG_DATA 0xCFC

static u32 pci_read(u8 bus, u8 slot, u8 func, u8 offset) {
    u32 addr = (1u << 31) | ((u32)bus << 16) | ((u32)slot << 11) |
               ((u32)func << 8) | (offset & 0xFC);
    outd(PCI_CONFIG_ADDR, addr);
    return ind(PCI_CONFIG_DATA);
}

static void pci_write(u8 bus, u8 slot, u8 func, u8 offset, u32 value) {
    u32 addr = (1u << 31) | ((u32)bus << 16) | ((u32)slot << 11) |
               ((u32)func << 8) | (offset & 0xFC);
    outd(PCI_CONFIG_ADDR, addr);
    outd(PCI_CONFIG_DATA, value);
}

typedef struct {
    u32 dba;
    u32 dbau;
    u32 reserved;
    u32 dbc;        // byte count - 1, bit 31 = interrupt on completion
} hba_prdt_t;

#define AHCI_PRDT_PER_TABLE 8

typedef struct {
    u8         cfis[64];
    u8         acmd[16];
    u8         reserved[48];
    hba_prdt_t prdt[AHCI_PRDT_PER_TABLE];
} hba_cmd_table_t;

// hba_cmd_table_t size = 64 + 16 + 48 + (8 * 16) = 256 bytes
// two fit per 4KiB page
typedef struct {
    u16 flags;      // [4:0] FIS len in DWORDs, [6] write
    u16 prdtl;
    u32 prdbc;
    u32 ctba;
    u32 ctbau;
    u32 reserved[4];
} hba_cmd_header_t;

typedef struct {
    u8 dsfis[28];
    u8 pad0[4];
    u8 psfis[20];
    u8 pad1[12];
    u8 rfis[20];
    u8 pad2[4];
    u8 sdbfis[8];
    u8 ufis[64];
    u8 reserved[96];
} hba_fis_t;

// sizeof(hba_fis_t) = 256 bytes

typedef volatile struct {
    u32 clb, clbu;
    u32 fb, fbu;
    u32 is;
    u32 ie;
    u32 cmd;
    u32 reserved0;
    u32 tfd;
    u32 sig;
    u32 ssts;
    u32 sctl;
    u32 serr;
    u32 sact;
    u32 ci;
    u32 sntf;
    u32 fbs;
    u32 reserved1[11];
    u32 vendor[4];
} hba_port_t;

typedef volatile struct {
    u32 cap;
    u32 ghc;
    u32 is;
    u32 pi;
    u32 vs;
    u32 ccc_ctl;
    u32 ccc_pts;
    u32 em_loc;
    u32 em_ctl;
    u32 cap2;
    u32 bohc;
    u8  reserved[0xA0 - 0x2C];
    u8  vendor[0x100 - 0xA0];
    hba_port_t ports[32];
} hba_mem_t;

typedef struct {
    u8  fis_type;   // 0x27
    u8  c;          // bit 7 = command
    u8  command;
    u8  featurel;
    u8  lba0, lba1, lba2;
    u8  device;
    u8  lba3, lba4, lba5;
    u8  featureh;
    u8  countl, counth;
    u8  icc;
    u8  control;
    u8  reserved[4];
} fis_reg_h2d_t;

#define ATA_CMD_READ_DMA_EXT  0x25
#define ATA_CMD_WRITE_DMA_EXT 0x35
#define ATA_CMD_IDENTIFY      0xEC
#define ATA_DEV_BUSY          0x80
#define ATA_DEV_DRQ           0x08


#define HBA_GHC_AE    (1u << 31)
#define HBA_CMD_ST    (1u << 0)
#define HBA_CMD_FRE   (1u << 4)
#define HBA_CMD_FR    (1u << 14)
#define HBA_CMD_CR    (1u << 15)
#define HBA_IS_TFES   (1u << 30)   // task file error status


#define AHCI_MAX_PORTS   32
#define AHCI_CMD_SLOTS   32

typedef struct {
    hba_cmd_header_t *cmd_list;                     // virtual, 1 page (1KB used)
    hba_fis_t        *fis;                          // virtual, fits after cmd_list
    hba_cmd_table_t  *cmd_tables[AHCI_CMD_SLOTS];   // virtual pointers
    int               active;
    u64               sector_count;
    u64               sector_size;
} ahci_port_t;

static hba_mem_t  *hba;
static ahci_port_t ports[AHCI_MAX_PORTS];

static void port_stop(hba_port_t *port) {
    port->cmd &= ~HBA_CMD_ST;
    port->cmd &= ~HBA_CMD_FRE;
    while (port->cmd & (HBA_CMD_CR | HBA_CMD_FR))
        ;
}

static void port_start(hba_port_t *port) {
    while (port->cmd & HBA_CMD_CR)
        ;
    port->cmd |= HBA_CMD_FRE;
    port->cmd |= HBA_CMD_ST;
}

static int find_free_slot(hba_port_t *port) {
    u32 slots = port->sact | port->ci;
    for (int i = 0; i < AHCI_CMD_SLOTS; i++)
        if (!(slots & (1u << i)))
            return i;
    return -1;
}

static void port_init(int idx) {
    hba_port_t  *port = &hba->ports[idx];
    ahci_port_t *p    = &ports[idx];

    port_stop(port);

    //
    // Page 0: command list (1024 bytes) + FIS buffer (256 bytes)
    //
    //   [0   .. 1023] : hba_cmd_header_t[32]   (command list)
    //   [1024.. 1279] : hba_fis_t               (received FIS)
    //
    // Both fit in one 4KiB page. FIS is 256-byte aligned because 1024
    // is a multiple of 256.
    //
    u8 *page0 = (u8 *)alloc(4096);
    memset(page0, 0, 4096);

    p->cmd_list = (hba_cmd_header_t *)page0;
    p->fis      = (hba_fis_t *)(page0 + 1024);

    u64 phys0        = virt_to_phys(page0);
    u64 phys_fis     = phys0 + 1024;

    port->clb  = (u32)(phys0);
    port->clbu = (u32)(phys0 >> 32);
    port->fb   = (u32)(phys_fis);
    port->fbu  = (u32)(phys_fis >> 32);

    //
    // Command tables: 256 bytes each, two per 4KiB page.
    // 32 slots → 16 pages.
    //
    for (int i = 0; i < AHCI_CMD_SLOTS; i += 2) {
        u8 *page = (u8 *)alloc(4096);
        memset(page, 0, 4096);

        u64 phys = virt_to_phys(page);

        for (int j = 0; j < 2 && (i + j) < AHCI_CMD_SLOTS; j++) {
            int slot = i + j;
            hba_cmd_table_t *tbl = (hba_cmd_table_t *)(page + j * 256);
            p->cmd_tables[slot] = tbl;

            u64 tbl_phys = phys + j * 256;
            p->cmd_list[slot].ctba  = (u32)(tbl_phys);
            p->cmd_list[slot].ctbau = (u32)(tbl_phys >> 32);
            p->cmd_list[slot].prdtl = AHCI_PRDT_PER_TABLE;
        }
    }

    // Clear pending errors and interrupts
    port->serr = 0xFFFFFFFF;
    port->is   = 0xFFFFFFFF;

    // Enable D2H register FIS interrupt + task file error
    port->ie = (1u << 0) | HBA_IS_TFES;

    port_start(port);
    p->active = 1;

    
    printf("AHCI: port %d initialized, sig=%x\n", idx, port->sig);
}

static void port_identify(int idx) {
    ahci_port_t *p    = &ports[idx];
    hba_port_t  *port = &hba->ports[idx];

    // IDENTIFY needs a 512-byte buffer; alloc gives us a page
    u16 *buf  = (u16 *)alloc(4096);
    u64  phys = virt_to_phys(buf);

    int slot = find_free_slot(port);
    if (slot < 0) return;

    hba_cmd_header_t *hdr = &p->cmd_list[slot];
    hdr->flags = sizeof(fis_reg_h2d_t) / 4;  // FIS length, no write bit
    hdr->prdtl = 1;
    hdr->prdbc = 0;

    hba_cmd_table_t *tbl = p->cmd_tables[slot];
    memset(tbl, 0, sizeof(hba_cmd_table_t));

    tbl->prdt[0].dba  = (u32)(phys);
    tbl->prdt[0].dbau = (u32)(phys >> 32);
    tbl->prdt[0].dbc  = 511; // 512 bytes - 1, no interrupt bit needed

    fis_reg_h2d_t *fis = (fis_reg_h2d_t *)tbl->cfis;
    fis->fis_type = 0x27;
    fis->c        = (1 << 7);
    fis->command  = ATA_CMD_IDENTIFY;
    fis->device   = 0;

    // Wait for idle
    int spin = 0;
    while ((port->tfd & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && ++spin < 1000000)
        ;

    port->ci = (1u << slot);

    while (port->ci & (1u << slot)) {
        if (port->is & HBA_IS_TFES) {
            printf("AHCI: identify error on port %d\n", idx);
            return;
        }
    }

    // Word 100-103: 48-bit LBA sector count
    p->sector_count = ((u64)buf[103] << 48) | ((u64)buf[102] << 32) |
                      ((u64)buf[101] << 16) | ((u64)buf[100]);


    // word 106 contains sector size info
    u16 w106 = buf[106];

    if (w106 & (1 << 14) && !(w106 & (1 << 15))) {
        // words 117-118: logical sector size in words (if bit 12 of w106 is set)
        if (w106 & (1 << 12)) {
            u32 logical_size = ((u32)buf[118] << 16) | buf[117];
            logical_size *= 2; // words to bytes
            p->sector_size = logical_size;
        }
    }

    printf("AHCI: port %d sections = %d\n", idx, p->sector_count);
}

// buf must be physically contiguous. For >AHCI_PRDT_PER_TABLE*4MB you'd need
// to split the request; at 8 PRDT entries of up to 4MB each that's 32MB max
// per call which is plenty for now.
static int ahci_rw(int idx, u64 lba, u32 sector_count, void *buf, int write) {
    ahci_port_t *p    = &ports[idx];
    hba_port_t  *port = &hba->ports[idx];

    if (!p->active) {
        print("AHCI: port not active\n");
        return -1;
    }

    int slot = find_free_slot(port);
    if (slot < 0) {
        print("AHCI: no free command slots\n");
        return -1;
    }

    hba_cmd_header_t *hdr = &p->cmd_list[slot];
    hdr->flags = (u16)(sizeof(fis_reg_h2d_t) / 4);
    if (write) hdr->flags |= (1 << 6);
    hdr->prdbc = 0;

    hba_cmd_table_t *tbl = p->cmd_tables[slot];
    memset(tbl, 0, sizeof(hba_cmd_table_t));

    // Fill PRDT — each entry covers up to 4MB (0x3FFFFF bytes + 1).
    // We use 512KB chunks to stay conservative and keep entry count low.
    #define CHUNK_SIZE (512 * 1024)

    u32 remaining   = sector_count * 512;
    u64 buff_phys    = virt_to_phys(buf);
    u16 prdt_count  = 0;

    while (remaining > 0 && prdt_count < AHCI_PRDT_PER_TABLE) {
        u32 chunk = remaining < CHUNK_SIZE ? remaining : CHUNK_SIZE;

        tbl->prdt[prdt_count].dba  = (u32)(buff_phys);
        tbl->prdt[prdt_count].dbau = (u32)(buff_phys >> 32);
        // dbc is bytes - 1; set interrupt on completion only on last entry
        tbl->prdt[prdt_count].dbc  = (chunk - 1);

        buff_phys  += chunk;
        remaining -= chunk;
        prdt_count++;
    }

    // Interrupt on last PRDT entry
    tbl->prdt[prdt_count - 1].dbc |= (1u << 31);

    hdr->prdtl = prdt_count;

    // Build the command FIS
    fis_reg_h2d_t *fis = (fis_reg_h2d_t *)tbl->cfis;
    fis->fis_type = 0x27;
    fis->c        = (1 << 7);
    fis->command  = write ? ATA_CMD_WRITE_DMA_EXT : ATA_CMD_READ_DMA_EXT;
    fis->device   = (1 << 6); // LBA mode

    fis->lba0 = (u8)(lba >>  0);
    fis->lba1 = (u8)(lba >>  8);
    fis->lba2 = (u8)(lba >> 16);
    fis->lba3 = (u8)(lba >> 24);
    fis->lba4 = (u8)(lba >> 32);
    fis->lba5 = (u8)(lba >> 40);

    fis->countl = (u8)(sector_count >> 0);
    fis->counth = (u8)(sector_count >> 8);

    // Wait until port is not busy
    int spin = 0;
    while ((port->tfd & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && ++spin < 1000000)
        ;
    if (spin >= 1000000) {
        print("AHCI: port hung before issue\n");
        return -1;
    }

    port->is = 0xFFFFFFFF; // clear interrupts before issuing
    port->ci = (1u << slot);

    // Poll for completion
    while (1) {
        if (!(port->ci & (1u << slot)))
            break;
        if (port->is & HBA_IS_TFES) {
            print("AHCI: task file error (port ");
            print32(idx);
            print(")\n");
            return -1;
        }
    }

    // Final error check after CI clears
    if (port->is & HBA_IS_TFES) {
        print("AHCI: task file error after completion\n");
        return -1;
    }

    return 0;
}

int ahci_read(int port, u64 lba, u32 count, void *buf) {
    int ret = ahci_rw(port, lba, count, buf, 0);
    if(ret < 0) {
        printf("Failed ahci read (%d, 0x", port);
        print64(lba);
        printf(", %u, %u)\n", count, buf);
    }
    return ret;
}

int ahci_write(int port, u64 lba, u32 count, void *buf) {
    return ahci_rw(port, lba, count, buf, 1);
}

u64 ahci_sector_count(int port) {
    return ports[port].sector_count;
}

int ahci_port_active(int port) {
    return ports[port].active;
}

void ahci_init() {
    // Scan PCI for AHCI controller (class=0x01, subclass=0x06)
    u8  found_bus = 0, found_slot = 0, found_func = 0;
    int found = 0;

    for (u32 bus = 0; bus < 256 && !found; bus++) {
        for (u8 slot = 0; slot < 32 && !found; slot++) {
            for (u8 func = 0; func < 8 && !found; func++) {
                u32 id = pci_read(bus, slot, func, 0x00);
                if ((id & 0xFFFF) == 0xFFFF) continue;

                u32 class = pci_read(bus, slot, func, 0x08);
                if (((class >> 24) & 0xFF) == 0x01 &&
                    ((class >> 16) & 0xFF) == 0x06) {
                    found_bus  = bus;
                    found_slot = slot;
                    found_func = func;
                    found = 1;
                }
            }
        }
    }

    if (!found) {
        print("AHCI: no controller found\n");
        return;
    }

    // Enable bus mastering + memory space
    u32 pcicmd = pci_read(found_bus, found_slot, found_func, 0x04);
    pcicmd |= (1 << 1) | (1 << 2);
    pci_write(found_bus, found_slot, found_func, 0x04, pcicmd);

    // BAR5 = ABAR
    u32 bar5      = pci_read(found_bus, found_slot, found_func, 0x24);
    u64 abar_phys = (u64)(bar5 & 0xFFFFF000);

    hba = (hba_mem_t *)map_mmio(abar_phys, 1<<12);
    if (!hba) {
        print("AHCI: failed to map HBA\n");
        return;
    }

    printf("AHCI: HBA mapped, version=%d", hba->vs);

    // Take AHCI ownership from BIOS if BOHC is supported (CAP2.BOH)
    if (hba->cap2 & (1 << 0)) {
        hba->bohc |= (1 << 1); // set OOS
        while (hba->bohc & (1 << 0)) // wait for BOS to clear
            ;
    }

    // Enable AHCI mode
    hba->ghc |= HBA_GHC_AE;

    // Init all ports that have a device attached
    u32 pi = hba->pi;
    for (int i = 0; i < AHCI_MAX_PORTS; i++) {
        if (!(pi & (1u << i))) continue;

        u32 ssts = hba->ports[i].ssts;
        u8  det  = (ssts >> 0) & 0xF;  // device detection
        u8  ipm  = (ssts >> 8) & 0xF;  // interface power management

        if (det == 3 && ipm == 1) {
            port_init(i);
            port_identify(i);
        }
    }
}
