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

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

static struct atapi_drive a_drive = { 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 < 4000; ++i)
        //printf("wait\r\n");
        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 };
        memcpy(&a_drive, &d, sizeof(struct atapi_drive));

        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);
    if (is_atapi_drive(PRIMARY_REG, ATA_PORT_MASTER))
        return;

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

    /* 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);
    if (is_atapi_drive(SECONDARY_REG, ATA_PORT_MASTER))
        return;

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

/* IO functions */
static 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)
{
    select_drive(a_drive.reg, a_drive.drive);
    wait_device_selection(a_drive.reg);

    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, a_drive.reg, CD_BLOCK_SZ);

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

    uint8_t read = 0;
    while (read != PACKET_COMMAND_COMPLETE)
        read = inb(ATA_REG_SECTOR_COUNT(a_drive.reg));

    return block;
}