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

#define MAX_IDE_DRIVES 8

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

struct ide_drive_info ide_drives[MAX_IDE_DRIVES];
drive_id_t n_ide_drives = 0;

void spin_delay() {
  for (uint32_t i = -1; i; --i)
    ;
}

//return true on error condition
bool poll(struct ide_drive_info *info) {
  spin_delay();

  uint8_t status;

  while ((status = inb(info->base_port | ATA_REG_STATUS)) & ATA_STATUS_BUSY)
    ;

  return (status & (ATA_STATUS_ERROR | ATA_STATUS_DRIVE_FAULT)) || !(status & ATA_STATUS_DRIVE_READY);
}

uint8_t ide_ata_rs(struct drive *d, uint32_t start, uint32_t count, void *buffer) {
  if (start & 0xf0000000)
    panic("IDE ATA driver only supports LBA28 addressing currently.");

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

  if (start + count > d->n_sectors)
    count = d->n_sectors - start;

  if (count & 0xffffff00)
    panic("IDE ATA driver only supports reading up to 128k at a time currently.");

  uint32_t lba = start & 0x00ffffff;

  uint32_t spin = -1;
  while (inb(info->base_port | ATA_REG_STATUS) & ATA_STATUS_BUSY)
    if (!spin--)
      panic("Spun out in IDE ATA reading");

  outb(info->base_port | ATA_REG_SELECT, (info->slave ? 0xf0 : 0xe0) | (start >> 24));
  outb(info->base_port | ATA_REG_COUNT_LOW, 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);

  uint16_t *buffer_16 = buffer;

  for (uint16_t i = 0; i < count; ++i) {
    if (poll(info))
      return i;
    for (uint16_t j = 0; j < 256; ++j)
      *buffer_16++ = inw(info->base_port + ATA_REG_DATA);
  }

  return count;
}

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

uint8_t ide_atapi_rs(struct drive *d, uint32_t start, uint32_t count, void *buffer) {
  //panic("IDE ATAPI reading not implemented yet");
  return 0;
}

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

struct id_space {
  uint16_t device_type;//0
  uint8_t skip1[98 - 2];
  uint16_t capabilities;//98
  uint8_t skip2[120 - (98 + 2)];
  uint32_t max_lba;//120
  uint8_t skip3[164 - (120 + 4)];
  uint32_t command_sets;//164
  uint8_t skip4[200 - (164 + 4)];
  uint32_t max_lba_ext;//200
  uint8_t skip5[512 - (200 + 4)];
} __attribute__ ((__packed__));

void read_id_space(uint16_t base_port, struct id_space *to) {
  uint32_t *to_32 = (uint32_t *)to;
  for (uint8_t i = 0; i < 128; ++i)
    *(to_32++) = ind(base_port | ATA_REG_DATA);
}

void test_drive(uint16_t base_port, uint16_t control_port, bool slave) {
  outb(base_port | ATA_REG_SELECT, slave ? ATA_SLAVE_SELECT : ATA_MASTER_SELECT);
  spin_delay();

  outb(base_port | ATA_REG_CMD, ATA_CMD_ID);
  spin_delay();

  if (!inb(base_port | ATA_REG_STATUS))
    return;

  if (n_ide_drives == MAX_IDE_DRIVES)
    panic("Maximum IDE drives reached");

  outb(base_port | ATA_REG_CONTROL, ATA_CONTROL_NO_IRQS);

  struct ide_drive_info *this_drive = ide_drives + n_ide_drives;
  this_drive->base_port = base_port;
  this_drive->control_port = control_port;
  this_drive->slave = slave;

  struct drive data;
  data.drive_id = n_ide_drives;

  uint32_t spin_out = -1;
  while (1) {//wait for identify to complete
    uint8_t status = inb(base_port | ATA_REG_STATUS);
    if (status & ATA_STATUS_ERROR) {
      uint16_t type = inb(base_port | ATA_REG_LBA1) +
                     (inb(base_port | ATA_REG_LBA2) << 8);

      //These two are listed on OSDev Wiki
      //  as being ATAPI device types.
      if ((type != 0xeb14) && (type != 0x9669))
        return;

      outb(base_port | ATA_REG_CMD, ATAPI_CMD_ID);

      data.read_sectors = &ide_atapi_rs;
      data.write_sectors = &ide_atapi_ws;
      data.drive_type = "IDE ATAPI";
      break;
    }
    if (!(status & ATA_STATUS_BUSY) && (status & ATA_STATUS_DATA_READY)) {
      data.read_sectors = &ide_ata_rs;
      data.write_sectors = &ide_ata_ws;
      data.drive_type = "IDE ATA";
      break;
    }
    if (!spin_out--)
      panic("IDE ATA identify won't complete");
  }

  struct id_space id;
  read_id_space(base_port, &id);

  if (!(id.capabilities & ATA_CAP_LBA))
    //TODO: CHS compability driver?
    return;
  data.n_sectors = (id.command_sets & ATA_SET_EXT) ? id.max_lba_ext : id.max_lba;

  commit_drive(data);
  ++n_ide_drives;
}

void init_ide() {
  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;

    uint16_t primary_base_port      = device->bar0 <= 1 ? 0x01f0 : device->bar0;
    uint16_t primary_control_port   = device->bar1 <= 1 ? 0x03f6 : device->bar1;
    uint16_t secondary_base_port    = device->bar2 <= 1 ? 0x0170 : device->bar2;
    uint16_t secondary_control_port = device->bar3 <= 1 ? 0x0376 : device->bar3;

    test_drive(primary_base_port,   primary_control_port,   false);
    test_drive(primary_base_port,   primary_control_port,   true);
    test_drive(secondary_base_port, secondary_control_port, false);
    test_drive(secondary_base_port, secondary_control_port, true);
  }
}