#include <asm/port_io.h>
#include <kernel/hw/ata.hpp>
#include <kernel/stdio.h>
#include <kernel/syscall.hpp>
#include <kernel/tty.h>
#include <kernel/vfs.h>
#include <types/allocator.hpp>
#include <types/status.h>
#include <types/stdint.h>

hw::ata::ata(port_id_t p)
    : data(p)
    , error(p + 1)
    , feats(p + 1)
    , count(p + 2)
    , lbalo(p + 3)
    , lbami(p + 4)
    , lbahi(p + 5)
    , drive(p + 6)
    , stats(p + 7)
    , comms(p + 7)
    , slave_flag(0x00)
{
}

hw::ata::stat_t hw::ata::status(void) const
{
    return hw::ata::stat_t { *stats };
}

bool hw::ata::identify(void) const
{
    char buf[512] {};

    drive = 0xa0 | slave_flag;
    count = 0;
    lbalo = 0;
    lbami = 0;
    lbahi = 0;
    comms = 0xec;

    stat_t stat {};
    while ((stat = status()).in.bsy)
        ;

    if (stat.in.err)
        return false;

    read_data(buf, 512);

    if (!status().in.rdy)
        return false;

    return true;
}

int hw::ata::select(bool master)
{
    if (master)
        slave_flag = 0x00;
    else
        slave_flag = 0x10;

    drive = 0xe0 | slave_flag;
    return GB_OK;
}

size_t hw::ata::read_data(char* buf, size_t n) const
{
    size_t orig_n = n;
    n /= 2;
    while (status().in.drq && n--) {
        *(uint16_t*)buf = *data;
        buf += sizeof(uint16_t);
    }
    return orig_n - n * 2;
}

size_t hw::ata::write_data(const char* buf, size_t n) const
{
    size_t orig_n = n;
    n /= 2;
    while (status().in.drq && n--) {
        data = *(uint16_t*)buf;
        buf += sizeof(uint16_t);
    }
    return orig_n - n * 2;
}

int hw::ata::read_sector(char* buf, uint32_t lba_low, uint16_t lba_high) const
{
    count = 0x00; // HIGH BYTE
    lbalo = (lba_low >> 24) & 0xff;
    lbami = lba_high & 0xff;
    lbahi = (lba_high >> 8) & 0xff;
    count = 0x01; // LOW BYTE
    lbalo = lba_low & 0xff;
    lbami = (lba_low >> 8) & 0xff;
    lbahi = (lba_low >> 16) & 0xff;
    comms = 0x24; // READ SECTORS EXT

    while (status().in.bsy)
        ;
    if (status().in.drq)
        read_data(buf, 512);
    return GB_OK;
}

int hw::ata::write_sector(const char* buf, uint32_t lba_low, uint16_t lba_high) const
{
    count = 0x00; // HIGH BYTE
    lbalo = (lba_low >> 24) & 0xff;
    lbami = lba_high & 0xff;
    lbahi = (lba_high >> 8) & 0xff;
    count = 0x01; // LOW BYTE
    lbalo = lba_low & 0xff;
    lbami = (lba_low >> 8) & 0xff;
    lbahi = (lba_low >> 16) & 0xff;
    comms = 0x24; // READ SECTORS EXT

    while (status().in.bsy)
        ;
    if (status().in.drq)
        write_data(buf, 512);
    return GB_OK;
}

static hw::ata* ata_pri;

void hw::init_ata(void* data)
{
    if (data != nullptr)
        syscall(0x03);

    ata_pri = types::kernel_allocator_new<ata>(ATA_PRIMARY_BUS_BASE);
    ata_pri->identify();
    ata_pri->select(true);

    register_special_block(
        2, 0,
        [](char* buf, size_t buf_size, size_t offset, size_t n) -> size_t {
            // TODO: check buf_size
            char b[512] {};
            char* orig_buf = buf;
            size_t start = offset / 512;
            size_t end = (offset + n + 511) / 512;
            for (size_t i = start; i < end; ++i) {
                ata_pri->read_sector(b, i, 0);
                size_t to_copy = 0;
                if (offset)
                    to_copy = 512 - offset;
                else
                    to_copy = n > 512 ? 512 : n;
                memcpy(buf, b + offset, to_copy);
                offset = 0;
                buf += to_copy;
                n -= to_copy;
            }
            return buf - orig_buf;
        },
        [](const char* buf, size_t offset, size_t n) -> size_t {
            syscall(0x03);
            return n;
        });
    auto* hda = vfs_open("/dev/hda");
    char buf[512] {};
    vfs_read(hda, buf, 512, 1, 512);
    vfs_write(hda, buf, 1, 512);
}