#include <stdbool.h>
#include <stdint.h>
#include "drive.h"
#include "panic.h"
#include "util.h"
#include "ata.h"
#include "pci.h"

#define MAX_IDE_DRIVES 8

struct ide_drive_info {
  uint16_t base_port;
  uint16_t alt_port;
  bool slave;
};

static struct ide_drive_info ide_drives[MAX_IDE_DRIVES];
static drive_id_t n_ide_drives;

typedef uint16_t spinner_t;

//returns the status after waiting
static uint8_t wait_for_ready(uint16_t base_port) {
  for (spinner_t n = -1; n; --n) {
    uint8_t s = inb(base_port | ATA_REG_STATUS);
    if (s & ATA_STATUS_ERROR)
      PANIC("Error status in IDE driver.");
    if (s & ATA_STATUS_DRIVE_READY)
      return s;
  }
  PANIC("Spun out in IDE driver.");
}

//returns the status after waiting
static uint8_t wait_for_error_or_ready(uint16_t base_port) {
  for (spinner_t n = -1; n; --n) {
    uint8_t s = inb(base_port | ATA_REG_STATUS);
    if (s & (ATA_STATUS_DRIVE_READY | ATA_STATUS_ERROR))
      return s;
  }
  PANIC("Spun out in IDE driver.");
}

//returns the status after waiting
static uint8_t wait_for_data_ready_not_busy(uint16_t base_port) {
  for (spinner_t n = -1; n; --n) {
    uint8_t s = inb(base_port | ATA_REG_STATUS);
    if (!(s & ATA_STATUS_BUSY) && (s & ATA_STATUS_DATA_READY))
      return s;
    if (s & ATA_STATUS_ERROR)
      PANIC("Error status in IDE driver.");
  }
  PANIC("Spun out in IDE driver.");
}

static uint32_t ide_ata_rs(const struct drive *d, uint32_t start, uint32_t count, void *buffer) {
  if (start >= d->n_sectors)
    return 0;
  if (start + count > d->n_sectors)
    count = d->n_sectors - start;
  if (!count)
    return 0;

  if (start & 0xf0000000)
    PANIC("IDE ATA driver does not support reads starting past 256MiB currently.");
  if (count & 0xffffff00)
    PANIC("IDE ATA driver does not support reads over 128kiB in length currently.");

  uint32_t lba = start & 0x00ffffff;

  struct ide_drive_info *info = ide_drives + d->drive_id;

  outb(info->base_port | ATA_REG_SELECT, (info->slave ? 0xf0 : 0xe0) | (start >> 24));
  //spin?

  outb(info->base_port | ATA_REG_COUNT, count);
  outb(info->base_port | ATA_REG_LBA0, lba);
  outb(info->base_port | ATA_REG_LBA1, lba >> 8);
  outb(info->base_port | ATA_REG_LBA2, lba >> 16);

  outb(info->base_port | ATA_REG_CMD, ATA_CMD_READ);

  uint16_t *buf16 = (uint16_t *)buffer;

  for (uint16_t i = 0; i < count; ++i) {
    wait_for_data_ready_not_busy(info->base_port);
    for (uint16_t j = 0; j < 256; ++j)
      *buf16++ = inw(info->base_port | ATA_REG_DATA);
  }

  return count;
}

static uint32_t ide_ata_ws(const struct drive *d, uint32_t start, uint32_t count, const void *buffer) {
  PANIC("IDE ATA writing not implemented yet.");
  return 0;
}

static uint32_t ide_atapi_rs(const struct drive *d, uint32_t start, uint32_t count, void *buffer) {
  if (start >= d->n_sectors)
    return 0;
  if (start + count > d->n_sectors)
    count = d->n_sectors - start;
  if (!count)
    return 0;

  PANIC("IDE ATAPI reading not implemented yet.");
}

static uint32_t ide_atapi_ws(const struct drive *d, uint32_t start, uint32_t count, const void *buffer) {
  PANIC("IDE ATAPI writing not implemented yet.");
  return 0;
}

static void nop(const struct drive *d) { }

struct id_space {
  uint8_t skip1[120];
  uint32_t max_lba;//120
  uint8_t skip2[512 - (120 + 4)];
} __attribute__ ((__packed__ ));

static void test_drive(uint16_t base_port, uint16_t alt_port, bool slave) {
  if (n_ide_drives == MAX_IDE_DRIVES)
    PANIC("Maximum number of IDE drives reached.");

  struct ide_drive_info *next = ide_drives + n_ide_drives;
  struct drive next_d;
  next_d.drive_id = n_ide_drives;

  next->base_port = base_port;
  next->alt_port = alt_port;
  next->slave = slave;

  outb(base_port | ATA_REG_SELECT, slave ? 0xb0 : 0xa0);
  outb(base_port | ATA_REG_COUNT, 0);
  outb(base_port | ATA_REG_LBA0, 0);
  outb(base_port | ATA_REG_LBA1, 0);
  outb(base_port | ATA_REG_LBA2, 0);

  outb(base_port | ATA_REG_CMD, ATA_CMD_ID);

  uint8_t status = inb(base_port | ATA_REG_STATUS);
  if (!status || (status == 0x7f))
    return;

  next_d.ready = &nop;
  next_d.done = &nop;

  struct id_space ids;

  if (!(wait_for_error_or_ready(base_port) & ATA_STATUS_ERROR)) {
    wait_for_ready(base_port);
    for (uint16_t *ids_16 = (uint16_t *)&ids; ids_16 < (uint16_t *)&ids + 256; ++ids_16)
      *ids_16 = inw(base_port | ATA_REG_DATA);

    next_d.drive_type = "IDE PATA";
    next_d.read_sectors = &ide_ata_rs;
    next_d.write_sectors = &ide_ata_ws;
  }
  else {
    uint16_t code = inb(base_port | 0x4) + (inb(base_port | 0x5) << 8);
    if (!code) {
      PANIC("PATA identification aborted.");
    }
    if (code == 0xeb14) {
      next_d.drive_type = "IDE PATAPI";
      next_d.read_sectors = &ide_atapi_rs;
      next_d.write_sectors = &ide_atapi_ws;

      //TODO: issue proper identify command
      return;
    }
    else if (code == 0xc33c) {
      next_d.drive_type = "IDE SATA";
      
      //TODO: set up drive hooks, issue proper identify command
      return;
    }
  }

  //in the future, maybe add support for 48-bit LBA, and/or CHS addressing
  if (!ids.max_lba)
    PANIC("Encountered ATA drive that doesn't support 28-bit LBA.");
  next_d.n_sectors = ids.max_lba;

  commit_drive(next_d);
  ++n_ide_drives;
}

void init_ide() {
  n_ide_drives = 0;

  uint16_t check_from = 0;
  struct pci_device *device;
  //double parentheses to let gcc know this assignment isn't a mistake
  while ((device = find_pci_device_from_class_and_subclass(PCI_MASS_STORAGE, PCI_IDE, check_from, &check_from))) {
    ++check_from;

    test_drive(0x01f0, 0x03f6, false);
    test_drive(0x01f0, 0x03f6, true);
    test_drive(0x0170, 0x0376, false);
    test_drive(0x0170, 0x0376, true);
  }
}