#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include "atasd.h"
#include "storedev.h"
#include "hal/hal.h"
#include "util/util.h"
#include "kprintf.h"
#include "errors.h"
#include "hshtb.h"
#include "vfs/vfs.h"

#define ATA_REG_DATA      0x00
#define ATA_REG_ERROR     0x01
#define ATA_REG_SECCOUNT0 0x02
#define ATA_REG_LBA0      0x03
#define ATA_REG_LBA1      0x04
#define ATA_REG_LBA2      0x05
#define ATA_REG_DRIVE     0x06
#define ATA_REG_STATUS    0x07
#define ATA_REG_COMMAND   0x07

#define ATA_BSY   0x80
#define ATA_DRDY  0x40
#define ATA_DF    0x20
#define ATA_ERR   0x01
#define ATA_DRQ   0x08

#define ATA_CMD_READ_PIO_EXT  0x24
#define ATA_CMD_WRITE_PIO_EXT 0x34
#define ATA_CMD_CACHE_FLUSH_EXT 0xEA
#define ATA_CMD_IDENTIFY      0xEC

#define ATA_MASTER  0x00
#define ATA_SLAVE   0x01

int ata_wait(uint16_t iobase, uint32_t timeout, bool require_drq, int checkerr) {
  uint32_t i = 0;
  uint8_t st;
  for(;;) {
    st = io_in8(iobase + ATA_REG_STATUS);
    if (!(st & ATA_BSY))
      break;
    if (++i >= timeout)
      return -1;
  }
  if (require_drq) {
    i = 0;
    while (!(st & ATA_DRQ)) {
      if (st & ATA_ERR)
        return -1;
      st = io_in8(iobase + ATA_REG_STATUS);
      if (++i >= timeout)
        return -1;
    }
  }
  if (checkerr && (st & (ATA_DF | ATA_ERR)))
    return -1;
  return 0;
}

void ata_delay400ns(uint16_t ctrlbase) {
  io_in8(ctrlbase);
  io_in8(ctrlbase);
  io_in8(ctrlbase);
  io_in8(ctrlbase);
}

// caller must synchronize
uint64_t ata_probesize_bytes(uint16_t iobase, uint16_t ctrlbase, int devno) {
  uint16_t identifydata[0x100];

  io_out8(iobase + ATA_REG_DRIVE, 0xA0 | (devno << 4));
  ata_delay400ns(ctrlbase);

  io_out8(iobase + ATA_REG_SECCOUNT0, 0);
  io_out8(iobase + ATA_REG_LBA0, 0);
  io_out8(iobase + ATA_REG_LBA1, 0);
  io_out8(iobase + ATA_REG_LBA2, 0);

  io_out8(iobase + ATA_REG_COMMAND, ATA_CMD_IDENTIFY);

  if (ata_wait(iobase, 100000, false, 0) < 0)
    return 0;

  uint8_t st = io_in8(iobase + ATA_REG_STATUS);
  if (st == 0)
    return 0;

  if (ata_wait(iobase, 100000, true, 1) < 0)
    return 0;

  io_ins16(iobase + ATA_REG_DATA, identifydata, 256);

  bool lba48 = (identifydata[83] & (1<<10)) != 0;
  if (lba48) {
    uint64_t sectors = (uint64_t)identifydata[100]
                    |  ((uint64_t)identifydata[101] << 16)
                    |  ((uint64_t)identifydata[102] << 32)
                    |  ((uint64_t)identifydata[103] << 48);
    return sectors * (uint64_t)STOREDEV_ATASD_SECTORSIZE;
  } else {
    uint64_t sectors = (uint64_t)identifydata[60]
                    |  ((uint64_t)identifydata[61] << 16);
    return sectors * (uint64_t)STOREDEV_ATASD_SECTORSIZE;
  }
}

void ata_probe(void) {
  uint64_t probesize;

  probesize = ata_probesize_bytes(0x1F0, 0x3F6, ATA_MASTER);
  if (probesize > 0) {

    char hs[20];
    LOG("ata", "found ATA primary bus master, size = %s\n", human_size(probesize, hs, sizeof(hs)));

    AtaSdInitExtra extra = {
      .devno = ATA_MASTER,
      .capacity = probesize,
      .iobase = 0x1F0,
      .ctrlbase = 0x3F6,
    };
    storedev_create(STOREDEV_ATASD, "atasd-ch0-M", (void *)&extra);
  }

  probesize = ata_probesize_bytes(0x1F0, 0x3F6, ATA_SLAVE);
  if (probesize > 0) {
    char hs[20];
    LOG("ata", "found ATA primary bus slave, size = %s\n", human_size(probesize, hs, sizeof(hs)));

    AtaSdInitExtra extra = {
      .devno = ATA_SLAVE,
      .capacity = probesize,
      .iobase = 0x1F0,
      .ctrlbase = 0x3F6,
    };
    storedev_create(STOREDEV_ATASD, "atasd-ch0-S", (void *)&extra);
  }
  
  probesize = ata_probesize_bytes(0x170, 0x376, ATA_MASTER);
  if (probesize > 0) {
    char hs[20];
    LOG("ata", "found ATA secondary bus master, size = %s\n", human_size(probesize, hs, sizeof(hs)));

    AtaSdInitExtra extra = {
      .devno = ATA_MASTER,
      .capacity = probesize,
      .iobase = 0x170,
      .ctrlbase = 0x376,
    };
    storedev_create(STOREDEV_ATASD, "atasd-ch1-M", (void *)&extra);
  }

  probesize = ata_probesize_bytes(0x170, 0x376, ATA_SLAVE);
  if (probesize > 0) {
    char hs[20];
    LOG("ata", "found ATA secondary bus slave, size = %s\n", human_size(probesize, hs, sizeof(hs)));

    AtaSdInitExtra extra = {
      .devno = ATA_SLAVE,
      .capacity = probesize,
      .iobase = 0x170,
      .ctrlbase = 0x376,
    };
    storedev_create(STOREDEV_ATASD, "atasd-ch1-S", (void *)&extra);
  }
}

void ata_setup(uint16_t iobase, uint16_t ctrlbase, int devno, uint16_t sectors, uint64_t lba) {
  io_out8(ctrlbase, 0x02);
  
  io_out8(iobase + ATA_REG_DRIVE, 0x40 | (devno << 4));
  ata_delay400ns(ctrlbase);

  io_out8(iobase + ATA_REG_SECCOUNT0, (sectors >> 8) & 0xFF);
  io_out8(iobase + ATA_REG_LBA0, (lba >> 24) & 0xFF);
  io_out8(iobase + ATA_REG_LBA1, (lba >> 32) & 0xFF);
  io_out8(iobase + ATA_REG_LBA2, (lba >> 40) & 0xFF);

  io_out8(iobase + ATA_REG_SECCOUNT0, sectors & 0xFF);
  io_out8(iobase + ATA_REG_LBA0, (uint8_t)(lba & 0xFF));
  io_out8(iobase + ATA_REG_LBA1, (uint8_t)((lba >> 8) & 0xFF));
  io_out8(iobase + ATA_REG_LBA2, (uint8_t)((lba >> 16) & 0xFF));

}

int32_t ata_read(uint16_t iobase, uint16_t ctrlbase, int devno,
                 uint64_t sectorsize, uint8_t *const buffer,
                 ptrdiff_t sector, ptrdiff_t off, size_t size) {
  int32_t ret = E_OK;

  uint64_t lba = (uint64_t)(sector * (ptrdiff_t)sectorsize + off) / sectorsize;
  size_t sectors = size / sectorsize;
  if (size % sectorsize) sectors++;

  ata_setup(iobase, ctrlbase, devno, sectors, lba);
  io_out8(iobase + ATA_REG_COMMAND, ATA_CMD_READ_PIO_EXT);

  for (size_t s = 0; s < sectors; s++) {
    if (ata_wait(iobase, 100000, true, 1) < 0) { ret = E_BADIO; goto done; }
    io_ins16(iobase + ATA_REG_DATA, (uint16_t *)(buffer + s * sectorsize), sectorsize / 2);
  }

done:
  return ret;
}

int32_t ata_write(uint16_t iobase, uint16_t ctrlbase, int devno,
                  uint64_t sectorsize, const uint8_t *const buffer,
                  ptrdiff_t sector, ptrdiff_t off, size_t size) {
  int32_t ret = E_OK;

  uint64_t lba = (uint64_t)(sector * (ptrdiff_t)sectorsize + off) / sectorsize;
  size_t sectors = size / sectorsize;
  if (size % sectorsize) sectors++;

  ata_setup(iobase, ctrlbase, devno, sectors, lba);
  io_out8(iobase + ATA_REG_COMMAND, ATA_CMD_WRITE_PIO_EXT);

  for (size_t s = 0; s < sectors; s++) {
    if (ata_wait(iobase, 100000, true, 1) < 0) { ret = E_BADIO; goto done; }
    io_outs16(iobase + ATA_REG_DATA, (uint16_t *)(buffer + s * sectorsize), sectorsize / 2);
  }

  io_out8(iobase + ATA_REG_COMMAND, ATA_CMD_CACHE_FLUSH_EXT);
  if (ata_wait(iobase, 100000, false, 1) < 0) { ret = E_BADIO; goto done; }

done:
  return ret;
}

int32_t atasd_init(struct StoreDev *sd, void *extra) {
  AtaSdInitExtra *e = (AtaSdInitExtra *)extra;
  sd->sd.atasd.devno    = e->devno;
  sd->sd.atasd.capacity = e->capacity;
  sd->sd.atasd.iobase   = e->iobase;
  sd->sd.atasd.ctrlbase = e->ctrlbase;
  return E_OK;
}

int32_t atasd_read(struct StoreDev *sd, uint8_t *const buffer, ptrdiff_t sector, ptrdiff_t off, size_t size) {
  int32_t ret;
  spinlock_acquire(&sd->spinlock);

  AtaSd *ata = &sd->sd.atasd;

  ret = ata_read(ata->iobase, ata->ctrlbase, ata->devno,
                 sd->sectorsize, buffer, sector, off, size);

  spinlock_release(&sd->spinlock);
  return ret;
}

int32_t atasd_write(struct StoreDev *sd, const uint8_t *const buffer, ptrdiff_t sector, ptrdiff_t off, size_t size) {
  int32_t ret;
  spinlock_acquire(&sd->spinlock);

  AtaSd *ata = &sd->sd.atasd;

  ret = ata_write(ata->iobase, ata->ctrlbase, ata->devno,
                  sd->sectorsize, buffer, sector, off, size);
  
  spinlock_release(&sd->spinlock);
  return ret;
}

int32_t atasd_cleanup(struct StoreDev *sd) {
  (void)sd;
  return E_OK;
}

size_t atasd_capacity(struct StoreDev *sd) {
  size_t c = 0;
  spinlock_acquire(&sd->spinlock);
  c = sd->sd.atasd.capacity;
  spinlock_release(&sd->spinlock);
  return c;
}
