#include <drv/ATA.hpp>

#include "print.hpp"

using namespace drv;

ATA::ATA(bool master, std::uint16_t portBase)
    : dataPort(portBase),
      errorPort(portBase + 0x1),
      sectorCountPort(portBase + 0x2),
      lbaLowPort(portBase + 0x3),
      lbaMidPort(portBase + 0x4),
      lbaHiPort(portBase + 0x5),
      devicePort(portBase + 0x6),
      commandPort(portBase + 0x7),
      statusPort(portBase + 0x7),
      controlPort(portBase + 0x206),
      devBit(master ? 0x00 : 0x10)
{
}

void
ATA::identify()
{
    waitWhileBusy();

    devicePort.write(0xe0 | devBit);
    controlPort.write(0);

    devicePort.write(0xe0);
    if (statusPort.read() == 0xff) {
        return;
    }

    devicePort.write(0xe0 | devBit);
    sectorCountPort.write(0);
    lbaLowPort.write(0);
    lbaMidPort.write(0);
    lbaHiPort.write(0);
    commandPort.write(0xec); // identify command

    if (statusPort.read() == 0x00) {
        return;
    }

    if (!waitUntilReady()) {
        kprint("ERROR");
        return;
    }

    for (int i = 0; i < 256; ++i) {
        const std::uint16_t data = dataPort.read();
        char text[3] = { };
        text[0] = (data >> 8) & 0xff;
        text[1] = data & 0xff;
        kprint(text);
    }
    kprint("\n");
}

void
ATA::read28(std::uint32_t sectorNum, span<std::uint8_t> buf)
{
    if (sectorNum > 0x0fffffff) {
        return;
    }

    waitWhileBusy();

    devicePort.write(0xe0 | devBit | ((sectorNum >> 24) & 0x0f));
    errorPort.write(0);
    sectorCountPort.write(1);
    lbaLowPort.write(sectorNum & 0xff);
    lbaMidPort.write((sectorNum >> 8) & 0xff);
    lbaHiPort.write((sectorNum >> 16) & 0xff);
    commandPort.write(0x20);

    if (!waitUntilReady()) {
        kprint("ATA ERROR");
        return;
    }

    for (int i = 0; i < buf.size(); i += 2) {
        const std::uint16_t wdata = dataPort.read();
        buf[i] = wdata & 0x00FF;
        if (i + 1 < buf.size()) {
            buf[i + 1] = (wdata >> 8) & 0xff;
        }
    }

    for (int i = buf.size() + buf.size()%2; i < 512; i += 2) {
        dataPort.read();
    }
}

void
ATA::write28(std::uint32_t sectorNum, span<const std::uint8_t> buf)
{
    if (sectorNum > 0x0fffffff) {
        return;
    }
    if (buf.size() > 512) {
        return;
    }

    waitWhileBusy();

    devicePort.write(0xe0 | devBit | ((sectorNum >> 24) & 0x0f));
    errorPort.write(0);
    sectorCountPort.write(1);
    lbaLowPort.write(sectorNum & 0xff);
    lbaMidPort.write((sectorNum >> 8) & 0xff);
    lbaHiPort.write((sectorNum >> 16) & 0xff);
    commandPort.write(0x30);

    kprint("Writing to ATA Drive: ");

    for (int i = 0; i < buf.size(); i += 2) {
        std::uint16_t wdata = buf[i];
        if (i + 1 < buf.size()) {
            wdata |= static_cast<std::uint16_t>(buf[i + 1]) << 8;
        }
        dataPort.write(wdata);

        char text[3] = { };
        text[0] = wdata & 0xff;
        text[1] = (wdata >> 8) & 0xff;
        kprint(text);
    }

    for (int i = buf.size() + (buf.size()%2); i < 512; i += 2) {
        dataPort.write(0x0000);
    }
}

void
ATA::flush()
{
    waitWhileBusy();

    devicePort.write(0xe0 | devBit);
    commandPort.write(0xe7);

    if (statusPort.read() == 0x00) {
        return;
    }

    if (!waitUntilReady()) {
        kprint("ERROR");
    }
}

void
ATA::waitWhileBusy()
{
    while (statusPort.read() & 0x80) {
        // Do nothing.
    }
}

bool
ATA::waitUntilReady()
{
    std::uint8_t status = statusPort.read();
    while ((status & 0x80) && !(status & 0x01)) {
        status = statusPort.read();
    }
    return !(status & 0x01);
}