#include "src/kernel/ata.h"
#include "src/kernel/panic.h"
#include "src/kernel/util.h"

#define BSY 0b10000000
#define DF 0b00100000
#define DRQ 0b00001000
#define ERR 0b00000001

#define READ_SECTORS 0x20
#define WRITE_SECTORS 0x30
#define FLUSH_CACHE 0xE7

#define SECTOR_WORDS 256

void ata_read_sectors(uint32_t index, uint8_t count, void *to) {
  while (inb(0x1F7) & BSY) {
  }

  outb(0x1F2, count);
  outb(0x1F3, index & 0xFF);
  index = index >> 8;
  outb(0x1F4, index & 0xFF);
  index = index >> 8;
  outb(0x1F5, index & 0xFF);
  index = index >> 8;
  outb(0x1F6, 0b11100000 | (index & 0b00001111));

  outb(0x1F7, READ_SECTORS);

  while (count--) {
    while (!(inb(0x1F7) & DRQ)) {
    }

    insw(0x1F0, to, SECTOR_WORDS);
    to = (uint16_t *)to + SECTOR_WORDS;
  }

  ASSERT(!(inb(0x1F7) & (ERR | DF)), "ata_read_sectors: read error");
}

void ata_write_sectors(uint32_t index, uint8_t count, const void *from) {
  while (inb(0x1F7) & BSY) {
  }

  outb(0x1F2, count);
  outb(0x1F3, index & 0xFF);
  index = index >> 8;
  outb(0x1F4, index & 0xFF);
  index = index >> 8;
  outb(0x1F5, index & 0xFF);
  index = index >> 8;
  outb(0x1F6, 0b11100000 | (index & 0b00001111));

  outb(0x1F7, WRITE_SECTORS);

  while (count--) {
    while (!(inb(0x1F7) & DRQ)) {
    }

    outsw(0x1F0, from, SECTOR_WORDS);
    from = (const uint16_t *)from + SECTOR_WORDS;
  }

  outb(0x1F7, FLUSH_CACHE);
  while (inb(0x1F7) & BSY) {
  }

  ASSERT(!(inb(0x1F7) & (ERR | DF)), "ata_write_sectors: write error");
}