#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      0x1f0
#define ATA_REG_ERROR     0x1f1
#define ATA_REG_SECCOUNT0 0x1f2
#define ATA_REG_LBA0      0x1f3
#define ATA_REG_LBA1      0x1f4
#define ATA_REG_LBA2      0x1f5
#define ATA_REG_DRIVE     0x1f6
#define ATA_REG_STATUS    0x1f7
#define ATA_REG_COMMAND   0x1f7
#define ATA_REG_CTRL      0x3f6

#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(uint32_t timeout, bool require_drq, int checkerr) {
  uint32_t i = 0;
  uint8_t st;
  for(;;) {
    st = io_in8(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(ATA_REG_STATUS);
      if (++i >= timeout)
        return -1;
    }
  }
  if (checkerr && (st & (ATA_DF | ATA_ERR)))
    return -1;
  return 0;
}

void ata_delay400ns(void) {
  io_in8(ATA_REG_CTRL);
  io_in8(ATA_REG_CTRL);
  io_in8(ATA_REG_CTRL);
  io_in8(ATA_REG_CTRL);
}

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

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

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

  io_out8(ATA_REG_COMMAND, ATA_CMD_IDENTIFY);

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

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

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

  io_ins16(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(int drive) {
  io_out8(ATA_REG_DRIVE, 0xA0 | (drive << 4));
  for (int i = 0; i < 1000; i++) {
    uint8_t status = io_in8(ATA_REG_STATUS);
    if (status != 0 && status != 0xFF)
      return true; // got something
  }
  return false;
}

void ata_probe(void) {
  if (ata_probe1(ATA_MASTER)) {
    uint64_t probesize = ata_probesize_bytes(ATA_MASTER);

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

    AtaSdInitExtra extra = {
      .devno = ATA_MASTER,
      .capacity = probesize,
    };
    StoreDev *sd = storedev_create(STOREDEV_ATASD, (void *)&extra);
  }
}

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

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

  io_out8(ATA_REG_SECCOUNT0, sectors & 0xFF);
  io_out8(ATA_REG_LBA0, (uint8_t)(lba & 0xFF));
  io_out8(ATA_REG_LBA1, (uint8_t)((lba >> 8) & 0xFF));
  io_out8(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;
  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->devno, sectors, lba);
  io_out8(ATA_REG_COMMAND, ATA_CMD_READ_PIO_EXT);

  for (size_t s = 0; s < sectors; s++) {
    if (ata_wait(100000, true, 1) < 0) { ret = E_BADIO; goto done; }
    io_ins16(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->devno, sectors, lba);
  io_out8(ATA_REG_COMMAND, ATA_CMD_WRITE_PIO_EXT);

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

  io_out8(ATA_REG_COMMAND, ATA_CMD_CACHE_FLUSH_EXT);
  if (ata_wait(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) {
  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;
}