#include <amd64/apic.h>
#include <amd64/intr_defs.h>
#include <amd64/io.h>
#include <device/device.h>
#include <device/idedrv.h>
#include <device/partitions.h>
#include <devices.h>
#include <irq/irq.h>
#include <libk/list.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/malloc.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <status.h>
#include <sys/debug.h>

#define IDE_REG_DATA     0x00
#define IDE_REG_ERROR    0x01
#define IDE_REG_SECCOUNT 0x02
#define IDE_REG_LBA0     0x03
#define IDE_REG_LBA1     0x04
#define IDE_REG_LBA2     0x05
#define IDE_REG_DRIVE    0x06
#define IDE_REG_STATUS   0x07
#define IDE_REG_CMD      0x07

#define IDE_BSY  0x80
#define IDE_DRDY 0x40
#define IDE_DF   0x20
#define IDE_ERR  0x01
#define IDE_DRQ  0x08

#define IDE_CMD_READ28   0x20
#define IDE_CMD_WRITE28  0x30
#define IDE_CMD_READ48   0x24
#define IDE_CMD_WRITE48  0x34
#define IDE_CMD_FLUSH48  0xEA
#define IDE_CMD_FLUSH28  0xE7
#define IDE_CMD_IDENTIFY 0xEC

#define IDE_READ  1
#define IDE_WRITE 2
#define IDE_FLUSH 3

static bool ide_wait (uint16_t io, uint32_t timeout, bool drq, bool errcheck) {
  uint32_t i = 0;
  uint8_t status;

  for (;;) {
    status = inb (io + IDE_REG_STATUS);

    if (!(status & IDE_BSY))
      break;

    if (++i >= timeout)
      return false;
  }

  if (drq) {
    i = 0;
    while (!(status & IDE_DRQ) || (status & IDE_BSY)) {
      if (status & IDE_ERR)
        return false;

      status = inb (io + IDE_REG_STATUS);

      if (status == 0xFF)
        return false;

      if (++i >= timeout)
        return false;
    }
  }

  if (errcheck && (status & (IDE_DF | IDE_ERR)))
    return false;

  return true;
}

#pragma clang optimize off
static void ide_delay (uint16_t ctrl) {
  inb (ctrl);
  inb (ctrl);
  inb (ctrl);
  inb (ctrl);
}
#pragma clang optimize on

static void ide_irq (void* arg, void* regs, bool user, struct reschedule_ctx* rctx) {
  uint64_t fd, fp;

  struct idedrv* idedrv = arg;

  spin_lock (&idedrv->device->lock, &fd);

  struct list_node_link* node = idedrv->requests;

  if (node == NULL) {
    (void)inb (idedrv->io + IDE_REG_STATUS);
    spin_unlock (&idedrv->device->lock, fd);
    return;
  }

  struct idedrv_request* req = list_entry (node, struct idedrv_request, requests_link);

  uint8_t status = inb (idedrv->io + IDE_REG_STATUS);

  if (req->type == IDE_FLUSH) {
    list_remove (idedrv->requests, &req->requests_link);
    spin_unlock (&idedrv->device->lock, fd);
    atomic_store (&req->done, 1);
    return;
  }

  if ((status & (IDE_ERR | IDE_DF))) {
    list_remove (idedrv->requests, &req->requests_link);
    spin_unlock (&idedrv->device->lock, fd);
    atomic_store (&req->done, 1);
    return;
  }

  if (!(status & IDE_BSY) && (status & IDE_DRQ)) {
    switch (req->type) {
    case IDE_READ: {
      uint16_t* p = req->buffer + (req->sector_done_count * (idedrv->sector_size / 2));
      insw (idedrv->io + IDE_REG_DATA, p, idedrv->sector_size / 2);
      req->sector_done_count++;
    } break;
    case IDE_WRITE: {
      req->sector_done_count++;
      if (req->sector_done_count < req->sector_count) {
        uint16_t* p = req->buffer + (req->sector_done_count * (idedrv->sector_size / 2));
        outsw (idedrv->io + IDE_REG_DATA, p, idedrv->sector_size / 2);
      }
    } break;
    }
  }

  if (req->sector_done_count >= req->sector_count) {
    list_remove (idedrv->requests, &req->requests_link);
    spin_unlock (&idedrv->device->lock, fd);
    atomic_store (&req->done, 1);
    return;
  }

  spin_unlock (&idedrv->device->lock, fd);
}

void ide_probe (uint16_t io, uint16_t ctrl, uint8_t devno, struct ide_probe* probe) {
  probe->flags = 0;
  probe->sector_count = 0;
  probe->sector_size = 0;

  uint16_t identify_buffer[256];

  uint8_t status = inb (io + IDE_REG_STATUS);

  if (status == 0xFF)
    return;

  outb (io + IDE_REG_DRIVE, 0xA0 | (devno << 4));
  ide_delay (ctrl);

  outb (io + IDE_REG_SECCOUNT, 0);
  outb (io + IDE_REG_LBA0, 0);
  outb (io + IDE_REG_LBA1, 0);
  outb (io + IDE_REG_LBA2, 0);

  outb (io + IDE_REG_CMD, IDE_CMD_IDENTIFY);

  status = inb (io + IDE_REG_STATUS);

  if (status == 0)
    return;

  if (!ide_wait (io, 100000, true, true)) {
    return;
  }

  insw (io + IDE_REG_DATA, identify_buffer, 256);

  probe->flags |= IDE_PROBE_AVAIL;

  if ((identify_buffer[106] & 0xC000) == 0x4000) {
    if (identify_buffer[106] & (1 << 12)) {
      uint32_t words_per_sector =
          (uint32_t)identify_buffer[117] | ((uint32_t)identify_buffer[118] << 16);
      probe->sector_size = (size_t)words_per_sector * 2;
    }
  }

  if ((identify_buffer[83] & (1 << 10)) != 0)
    probe->flags |= IDE_PROBE_LBA48;

  if ((probe->flags & IDE_PROBE_LBA48)) {
    probe->sector_count =
        (size_t)((uint64_t)identify_buffer[100] | ((uint64_t)identify_buffer[101] << 16) |
                 ((uint64_t)identify_buffer[102] << 32) | ((uint64_t)identify_buffer[103] << 48));
  } else {
    probe->sector_count =
        (size_t)((uint64_t)identify_buffer[60] | ((uint64_t)identify_buffer[61] << 16));
  }

  probe->io = io;
  probe->ctrl = ctrl;
  probe->devno = devno;

  if (probe->sector_size == 0)
    probe->sector_size = 512;
}

static void ide_prepare (struct idedrv* idedrv, size_t sector, uint16_t sector_count, bool clear) {
  if (clear)
    outb (idedrv->ctrl, 0x00);

  if (idedrv->lba48) {
    outb (idedrv->io + IDE_REG_DRIVE, 0x40 | (idedrv->devno << 4));
    ide_delay (idedrv->ctrl);

    outb (idedrv->io + IDE_REG_SECCOUNT, (sector_count >> 8) & 0xFF);
    outb (idedrv->io + IDE_REG_LBA0, (sector >> 24) & 0xFF);
    outb (idedrv->io + IDE_REG_LBA1, (sector >> 32) & 0xFF);
    outb (idedrv->io + IDE_REG_LBA2, (sector >> 40) & 0xFF);

    outb (idedrv->io + IDE_REG_SECCOUNT, sector_count & 0xFF);
    outb (idedrv->io + IDE_REG_LBA0, sector & 0xFF);
    outb (idedrv->io + IDE_REG_LBA1, (sector >> 8) & 0xFF);
    outb (idedrv->io + IDE_REG_LBA2, (sector >> 16) & 0xFF);
  } else {
    outb (idedrv->io + IDE_REG_DRIVE, 0xE0 | (idedrv->devno << 4) | ((sector >> 24) & 0xFF));
    ide_delay (idedrv->ctrl);

    uint8_t count = (sector_count == 256) ? 0 : (uint8_t)sector_count;
    outb (idedrv->io + IDE_REG_SECCOUNT, count);
    outb (idedrv->io + IDE_REG_LBA0, sector & 0xFF);
    outb (idedrv->io + IDE_REG_LBA1, (sector >> 8) & 0xFF);
    outb (idedrv->io + IDE_REG_LBA2, (sector >> 16) & 0xFF);
  }
}

bool idedrv_init (struct device* device, void* arg, struct proc* proc,
                  struct reschedule_ctx* rctx) {
  (void)proc, (void)rctx;

  struct idedrv_init* init = arg;

  struct idedrv* idedrv = malloc (sizeof (*idedrv));

  if (idedrv == NULL)
    return false;

  idedrv->device = device;
  idedrv->lba48 = init->lba48;
  idedrv->sector_count = init->sector_count;
  idedrv->sector_size = init->sector_size;
  idedrv->io = init->io;
  idedrv->ctrl = init->ctrl;
  idedrv->devno = init->devno;
  idedrv->primscnd = init->primscnd;

  device->udata = idedrv;

  if (idedrv->primscnd == 1) {
    ioapic_route_irq (IDE_DRIVE_PRIM, 14, 0, thiscpu->lapic_id);
    irq_attach (&ide_irq, idedrv, IDE_DRIVE_PRIM);
  } else if (idedrv->primscnd == 2) {
    ioapic_route_irq (IDE_DRIVE_SCND, 15, 0, thiscpu->lapic_id);
    irq_attach (&ide_irq, idedrv, IDE_DRIVE_SCND);
  }

  return true;
}

void idedrv_fini (struct device* device, struct proc* proc, struct reschedule_ctx* rctx) {
  struct idedrv* idedrv = device->udata;

  struct list_node_link *req_link, *tmp_req_link;
  list_foreach (idedrv->requests, req_link, tmp_req_link) {
    struct idedrv_request* req = list_entry (req_link, struct idedrv_request, requests_link);
    list_remove (idedrv->requests, &req->requests_link);
    free (req);
  }

  irq_detach (idedrv->primscnd == 1 ? IDE_DRIVE_PRIM : IDE_DRIVE_SCND);
  free (idedrv);
}

int idedrv_read (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                 uint64_t* lockflags, void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a4;

  if (a1 == NULL || a2 == NULL || a3 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t sector = *(size_t*)a1;
  size_t sector_count = *(size_t*)a2;
  uint16_t* buffer = a3;

  struct idedrv* idedrv = device->udata;

  if (sector + sector_count > idedrv->sector_count)
    return -ST_OOB_ERROR;

  struct idedrv_request* req = malloc (sizeof (*req));

  if (req == NULL)
    return -ST_OOM_ERROR;

  memset (req, 0, sizeof (*req));
  req->buffer = buffer;
  req->sector_count = sector_count;
  req->sector_done_count = 0;
  req->type = IDE_READ;

  list_append (idedrv->requests, &req->requests_link);

  ide_prepare (idedrv, sector, sector_count, true);

  uint8_t cmd = idedrv->lba48 ? IDE_CMD_READ48 : IDE_CMD_READ28;
  outb (idedrv->io + IDE_REG_CMD, cmd);

  if (!ide_wait (idedrv->io, 100000, true, true)) {
    list_remove (idedrv->requests, &req->requests_link);
    free (req);
    return -ST_XDRV_READ_ERROR;
  }

  insw (idedrv->io + IDE_REG_DATA, buffer, idedrv->sector_size / 2);

  req->sector_done_count = 1;

  spin_unlock (&device->lock, *lockflags);

  while (!atomic_load (&req->done))
    spin_lock_relax ();

  spin_lock (&device->lock, lockflags);

  free (req);

  return ST_OK;
}

int idedrv_write (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                  uint64_t* lockflags, void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a4;
  uint64_t fp;

  if (a1 == NULL || a2 == NULL || a3 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t sector = *(size_t*)a1;
  size_t sector_count = *(size_t*)a2;
  uint16_t* buffer = a3;

  struct idedrv* idedrv = device->udata;

  if (sector + sector_count > idedrv->sector_count)
    return -ST_OOB_ERROR;

  struct idedrv_request* req = malloc (sizeof (*req));

  struct idedrv_request* flushreq = malloc (sizeof (*req));

  if (req == NULL)
    return -ST_OOM_ERROR;

  if (flushreq == NULL) {
    free (req);
    return -ST_OOM_ERROR;
  }

  memset (req, 0, sizeof (*req));
  req->buffer = buffer;
  req->sector_count = sector_count;
  req->sector_done_count = 0;
  req->type = IDE_WRITE;

  list_append (idedrv->requests, &req->requests_link);

  memset (flushreq, 0, sizeof (*flushreq));
  flushreq->buffer = NULL;
  flushreq->sector_count = 0;
  flushreq->sector_done_count = 0;
  flushreq->type = IDE_FLUSH;

  list_append (idedrv->requests, &flushreq->requests_link);

  ide_prepare (idedrv, sector, sector_count, true);

  uint8_t cmd = idedrv->lba48 ? IDE_CMD_WRITE48 : IDE_CMD_WRITE28;
  outb (idedrv->io + IDE_REG_CMD, cmd);

  if (!ide_wait (idedrv->io, 100000, true, true)) {
    list_remove (idedrv->requests, &req->requests_link);
    free (req);
    return -ST_XDRV_WRITE_ERROR;
  }

  outsw (idedrv->io + IDE_REG_DATA, buffer, idedrv->sector_size / 2);

  req->sector_done_count = 1;

  spin_unlock (&device->lock, *lockflags);

  while (!atomic_load (&req->done))
    spin_lock_relax ();

  spin_lock (&device->lock, lockflags);

  if (idedrv->lba48)
    outb (idedrv->io + IDE_REG_CMD, IDE_CMD_FLUSH48);
  else
    outb (idedrv->io + IDE_REG_CMD, IDE_CMD_FLUSH28);

  spin_unlock (&device->lock, *lockflags);

  while (!atomic_load (&flushreq->done))
    spin_lock_relax ();

  spin_lock (&device->lock, lockflags);

  free (req);
  free (flushreq);

  return ST_OK;
}

int idedrv_get_device_type (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                            uint64_t* lockflags, void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)device, (void)a2, (void)a3, (void)a4;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  int* device_type = (int*)a1;

  *device_type = XDRV_TYPE_IDEDRV;

  return ST_OK;
}

int idedrv_get_sector_size (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                            uint64_t* lockflags, void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a2, (void)a3, (void)a4;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* secsize = (size_t*)a1;

  struct idedrv* idedrv = device->udata;

  *secsize = idedrv->sector_size;

  return ST_OK;
}

int idedrv_get_size (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                     uint64_t* lockflags, void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a2, (void)a3, (void)a4;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* size = (size_t*)a1;

  struct idedrv* idedrv = device->udata;

  *size = idedrv->sector_size * idedrv->sector_count;

  return ST_OK;
}