#include <k/atapi.h>
#include <string.h>

#include "atapi.h"
#include "io.h"

static struct atapi_drive drives[4] = { 0 };

/* wait functions */
static void busy_wait(uint16_t drive)
{
    int res = 0;
    while (res & BSY)
        res = inb(ATA_REG_STATUS(drive));
}

static void wait_device_selection(uint16_t drive)
{
    for (int i = 0; i < 4; ++i)
        inb(drive);
}

static void wait_packet_request(uint16_t drive)
{
    int res = 0;
    while (res & BSY && !(res & DRQ))
        res = inb(ATA_REG_STATUS(drive));
}

/* drive discover functions */
static void select_drive(uint16_t bus, uint8_t slave)
{
    outb(ATA_REG_DRIVE(bus), slave);
}

static uint8_t is_atapi_drive(uint16_t bus, uint8_t slave)
{
    static uint8_t atapi_sig[] = { ATAPI_SIG_SC, ATAPI_SIG_LBA_LO,
                                   ATAPI_SIG_LBA_MI, ATAPI_SIG_LBA_HI };
    uint8_t sig[4];

    sig[0] = inb(ATA_REG_SECTOR_COUNT(bus));
    sig[1] = inb(ATA_REG_LBA_LO(bus));
    sig[2] = inb(ATA_REG_LBA_MI(bus));
    sig[3] = inb(ATA_REG_LBA_HI(bus));

    if (!memcmp(sig, atapi_sig, 4))
    {
        struct atapi_drive d = { bus, slave };
        if (bus == PRIMARY_REG)
        {
            if (slave)
                drives[1] = d;
            else
                drives[0] = d;
        }
        else
        {
            if (slave)
                drives[3] = d;
            else
                drives[2] = d;
        }

        return 1;
    }

    return 0;
}

void discover_atapi_drive(void)
{
    /* primary_reg */
    outb(PRIMARY_DCR, SRST);
    outb(PRIMARY_DCR, INTERRUPT_DISABLE);

    // master_port
    select_drive(PRIMARY_REG, ATA_PORT_MASTER);
    wait_device_selection(PRIMARY_REG);
    is_atapi_drive(PRIMARY_REG, ATA_PORT_MASTER);

    // slave port
    select_drive(PRIMARY_REG, ATA_PORT_SLAVE);
    wait_device_selection(PRIMARY_REG);
    is_atapi_drive(PRIMARY_REG, ATA_PORT_SLAVE);

    /* secondary_reg */
    outb(SECONDARY_DCR, SRST);
    outb(SECONDARY_DCR, INTERRUPT_DISABLE);

    // master_port
    select_drive(SECONDARY_REG, ATA_PORT_MASTER);
    wait_device_selection(SECONDARY_REG);
    is_atapi_drive(SECONDARY_REG, ATA_PORT_MASTER);

    // slave port
    select_drive(SECONDARY_REG, ATA_PORT_SLAVE);
    wait_device_selection(SECONDARY_REG);
    is_atapi_drive(SECONDARY_REG, ATA_PORT_SLAVE);
}

/* IO functions */
int send_packet(struct SCSI_packet *pkt, uint16_t drive, uint16_t size)
{
    // cpy SCSI_packet into uint16_t array
    uint16_t packet[PACKET_SZ / 2];
    memcpy(packet, pkt, PACKET_SZ * sizeof(uint8_t));

    // wait for a packet to be requested
    busy_wait(drive);
    outb(ATA_REG_FEATURES(drive), 0);
    outb(ATA_REG_SECTOR_COUNT(drive), 0);
    outb(ATA_REG_LBA_MI(drive), size);
    outb(ATA_REG_LBA_HI(drive), size >> 8);
    outb(ATA_REG_COMMAND(drive), PACKET);
    wait_packet_request(drive);

    // write SCSI packet
    for (int i = 0; i < PACKET_SZ / 2; ++i)
        outw(ATA_REG_DATA(drive), packet[i]);
    uint8_t read = 0;
    while (read != PACKET_DATA_TRANSMIT)
        read = inb(ATA_REG_SECTOR_COUNT(drive));

    return 0;
}

static char block[CD_BLOCK_SZ];

void read_block(size_t lba, uint16_t drive)
{
    struct SCSI_packet pkt = { 0 };
    pkt.op_code  = READ_12;
    pkt.lba_lo   = lba & 0xff;
    pkt.lba_milo = (lba >> 0x8)  & 0xff;
    pkt.lba_mihi = (lba >> 0x10) & 0xff;
    pkt.lba_hi   = (lba >> 0x18) & 0xff;
    pkt.transfer_length_lo = 1;

    send_packet(&pkt, drive, PACKET_SZ);

    // we can now read
    uint16_t *buf = block;
    for (int i = 0; i < CD_BLOCK_SZ / 2; ++i)
        buf[i] = inw(ATA_REG_DATA(drive));

    uint8_t read = 0;
    while (read != PACKET_COMMAND_COMPLETE)
        read = inb(ATA_REG_SECTOR_COUNT(drive));
}