#include <stdint.h>
#include "panic.h"
#include "util.h"
#include "ata.h"

enum {
  AP_PRIMARY   = 0x01f0,
  AP_SECONDARY = 0x0170,

  AP_DATA  = 0x0,
  AP_ERROR = 0x1,
  AP_FEAT  = 0x1,
  AP_COUNT = 0x2,
  AP_SEC   = 0x3,
  AP_CYL_L = 0x4,
  AP_CYL_H = 0x5,
  AP_HD_DR = 0x6,
  AP_LBA0  = 0x3,
  AP_LBA1  = 0x4,
  AP_LBA2  = 0x5,
  AP_L3_DR = 0x6,
  AP_STAT  = 0x7,
  AP_CMD   = 0x7,

  APC_STAT = 0x0206,
  APC_CTRL = 0x0206,
  APC_ADDR = 0x0207
};

enum {
  ADR_BASE  = 0xa0,
  ADR_CASCD = 0x10,
  ADR_LBA   = 0x40
};

enum {
  AS_ERROR   = 0x01,
  AS_INDEX   = 0x02,
  AS_CORRECT = 0x04,
  AS_DREADY  = 0x08,
  AS_SERVICE = 0x10,
  AS_FAULT   = 0x20,
  AS_READY   = 0x40,
  AS_BUSY    = 0x80
};

enum {
  AC_READ  = 0x20,
  AC_WRITE = 0x30,
  AC_FLUSH = 0xe7
};

void ata_ldelay() {
  uint16_t i = 65535;
  while (i--)
    ;
}

void ata_sdelay() {
  uint8_t i = 255;
  while (i--)
    ;
}

void ata_dpoll() {
  uint8_t s;
  while ((s = inb(AP_PRIMARY | AP_STAT) & (AS_BUSY | AS_DREADY)) != AS_DREADY)
    if (s & (AS_ERROR | AS_FAULT))
      panic("ATA controller error or fault.");
}

uint8_t read_sectors(uint16_t start, uint8_t count, void *buffer) {
  outb(AP_PRIMARY | AP_L3_DR, ADR_BASE | ADR_LBA);
  outb(AP_PRIMARY | AP_LBA2, 0);
  outb(AP_PRIMARY | AP_LBA1, start >> 8);
  outb(AP_PRIMARY | AP_LBA0, start & 0xff);
  outb(AP_PRIMARY | AP_COUNT, count);
  outb(AP_PRIMARY | AP_CMD, AC_READ);
  uint8_t i = count - 1;
  uint16_t *ptr = buffer;
  do {
    ata_dpoll();
    uint8_t j = 255;
    do
      *(ptr++) = inw(AP_PRIMARY | AP_DATA);
    while (j--);
    ata_ldelay();
  } while (i--);
  return count;//TODO: return early if necessary
}

uint8_t write_sectors(uint16_t start, uint8_t count, void *buffer) {
  outb(AP_PRIMARY | AP_L3_DR, ADR_BASE | ADR_LBA);
  outb(AP_PRIMARY | AP_LBA2, 0);
  outb(AP_PRIMARY | AP_LBA1, start >> 8);
  outb(AP_PRIMARY | AP_LBA0, start & 0xff);
  outb(AP_PRIMARY | AP_COUNT, count);
  outb(AP_PRIMARY | AP_CMD, AC_WRITE);
  uint8_t i = 0;
  uint16_t *ptr = buffer;
  do {
    ata_dpoll();
    uint8_t j = 0;
    do {
      ata_sdelay();
      outw(AP_PRIMARY | AP_DATA, *(ptr++));
    }
    while (++j);
    ata_ldelay();
  } while (++i != count);
  outb(AP_PRIMARY | AP_CMD, AC_FLUSH);
  return count;//TODO: return early if necessary
}