#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;
  }
}

/* bool ata_probe1(uint16_t iobase, uint16_t ctrlbase, int drive) { */
/*   /1* io_out8(iobase + ATA_REG_DRIVE, 0xA0 | (drive << 4)); *1/ */
/*   /1* for (int i = 0; i < 1000; i++) { *1/ */
/*   /1*   uint8_t status = io_in8(iobase + ATA_REG_STATUS); *1/ */
/*   /1*   if (status != 0 && status != 0xFF) *1/ */
/*   /1*     return true; // got something *1/ */
/*   /1* } *1/ */
/*   /1* return false; *1/ */

/*   uint8_t cl, ch; */

/*   io_out8(iobase + ATA_REG_DRIVE, 0xA0 | (drive << 4)); */
/*   ata_delay400ns(ctrlbase); */

/*   io_out8(iobase + ATA_REG_SECCOUNT0, 0x55); */
/*   io_out8(iobase + ATA_REG_LBA0, 0xAA); */
/*   io_out8(iobase + ATA_REG_SECCOUNT0, 0xAA); */
/*   io_out8(iobase + ATA_REG_LBA0, 0x55); */

/*   cl = io_in8(iobase + ATA_REG_SECCOUNT0); */
/*   ch = io_in8(iobase + ATA_REG_LBA0); */

/*   if (cl != 0x55 || ch != 0xAA) */
/*     return false; */

/*   io_out8(iobase + ATA_REG_COMMAND, ATA_CMD_IDENTIFY); */
/*   ata_delay400ns(ctrlbase); */

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

/*   while ((st & (ATA_BSY | ATA_DRQ | ATA_ERR)) == ATA_BSY) */
/*     st = io_in8(iobase + ATA_REG_STATUS); */

/*   if (st & ATA_ERR) */
/*     return false; */

/*   if (!(st & ATA_DRQ)) */
/*     return false; */

/*   return true; */
/* } */

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, (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, (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, (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, (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 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;

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

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

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

  ret = E_OK;

done:
  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;

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

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

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

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

  ret = E_OK;
  
done:
  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;
}