#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 <device/pci.h>
#include <device/pci_info.h>
#include <devices.h>
#include <libk/fieldsizeof.h>
#include <libk/lengthof.h>
#include <libk/printf.h>
#include <libk/std.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <sys/debug.h>

static int ide_counter = 0;

static const char* progif_msg[] = {
    [0x00] = "ISA Compatibility mode-only controller",
    [0x05] = "PCI native mode-only controller",
    [0x0A] =
        "ISA Compatibility mode controller, supports both channels switched to PCI native mode",
    [0x0F] =
        "PCI native mode controller, supports both channels switched to ISA compatibility mode",
    [0x80] = "ISA Compatibility mode-only controller, supports bus mastering",
    [0x85] = "PCI native mode-only controller, supports bus mastering",
    [0x8A] =
        "ISA Compatibility mode controller, supports both channels switched to PCI native mode, supports bus mastering",
    [0x8F] =
        "PCI native mode controller, supports both channels switched to ISA compatibility mode, supports bus mastering",
};

static void ide_make_device (struct proc* proc, struct reschedule_ctx* rctx,
                             struct ide_probe probe) {
  DEBUG ("Found IDE drive: io=%x ctrl=%x no=%u sector=%zu count=%zu\n", probe.io, probe.ctrl,
         probe.devno, probe.sector_size, probe.sector_count);

  char device_key[fieldsizeof (struct device, key)];
  snprintf (device_key, sizeof (device_key), "IDE%d", ide_counter++);

  device_op_func_t ops[] = {
      [XDRV_GET_SIZE] = &idedrv_get_size,
      [XDRV_GET_SECTOR_SIZE] = &idedrv_get_sector_size,
      [XDRV_GET_DEVICE_TYPE] = &idedrv_get_device_type,
      [XDRV_READ] = &idedrv_read,
      [XDRV_WRITE] = &idedrv_write,
  };

  struct idedrv_init init = {
      .lba48 = ((probe.flags & IDE_PROBE_LBA48) != 0),
      .sector_count = probe.sector_count,
      .sector_size = probe.sector_size,
      .io = probe.io,
      .ctrl = probe.ctrl,
      .devno = probe.devno,
      .irq = probe.irq,
      .irqs_support = probe.irqs_support,
  };
  struct device* ide = device_create (device_key, ops, lengthof (ops), &idedrv_init, &idedrv_fini,
                                      &init, proc, rctx);
  device_probe_partitions (proc, rctx, ide);
}

bool pci_ide_init (struct proc* proc, struct reschedule_ctx* rctx, struct pci_info pci_info) {
  uint16_t pci_cmd = pci_read16 (pci_info.bus, pci_info.slot, pci_info.func, PCI_COMMAND);

  uint16_t new_cmd = (pci_cmd | (1 << 0) | (1 << 2)) & ~(1 << 10);

  if (pci_cmd != new_cmd) {
    pci_write16 (pci_info.bus, pci_info.slot, pci_info.func, PCI_COMMAND, new_cmd);
  }

  struct ide_probe probe;

  uint8_t progif = pci_read8 (pci_info.bus, pci_info.slot, pci_info.func, PCI_PROG_IF);
  DEBUG ("progif: %s\n", progif_msg[progif]);

  uint16_t pcmd, pctrl, scmd, sctrl;

  if ((progif & 0x01)) {
    uint32_t bar0 = pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR0);
    uint32_t bar1 = pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR1);

    if (!(bar0 & PCI_BAR_IO) || !(bar1 & PCI_BAR_IO)) {
      DEBUG ("Non IO BARs not supported\n");
      return false;
    }

    pcmd = (uint16_t)(bar0 & 0xFFFC);
    pctrl = (uint16_t)(bar1 & 0xFFFC);

    if (pctrl)
      pctrl += 2;
  } else {
    pcmd = 0x1F0;
    pctrl = 0x3F6;
  }

  if ((progif & 0x04)) {
    uint32_t bar2 = pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR2);
    uint32_t bar3 = pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR3);

    if (!(bar2 & PCI_BAR_IO) || !(bar3 & PCI_BAR_IO)) {
      DEBUG ("Non IO BARs not supported\n");
      return false;
    }

    scmd = (uint16_t)(bar2 & 0xFFFC);
    sctrl = (uint16_t)(bar3 & 0xFFFC);

    if (sctrl)
      sctrl += 2;
  } else {
    scmd = 0x170;
    sctrl = 0x376;
  }

  bool irqs_support = false;

  if ((progif & 0x05)) {
    irqs_support = false;
  } else {
    irqs_support = true;
    ioapic_route_irq (IDE_DRIVE_PRIM, 14, 0, thiscpu->lapic_id);
    ioapic_route_irq (IDE_DRIVE_SCND, 15, 0, thiscpu->lapic_id);
  }

  DEBUG ("pcmd=%x, pctrl=%x\n", pcmd, pctrl);
  DEBUG ("scmd=%x, sctrl=%x\n", scmd, sctrl);
  DEBUG ("IRQ support=%d\n", irqs_support);

  uint16_t channels[2][3] = {{pcmd, pctrl, IDE_DRIVE_PRIM}, {scmd, sctrl, IDE_DRIVE_SCND}};

  for (size_t i = 0; i < lengthof (channels); i++) {
    uint16_t cmd = channels[i][0];
    uint16_t ctrl = channels[i][1];
    uint8_t irq = channels[i][2];

    for (size_t dev = 0; dev < 2; dev++) {
      ide_probe (cmd, ctrl, dev, &probe);

      probe.ctrl = ctrl;
      probe.io = cmd;
      probe.irq = irq;
      probe.irqs_support = irqs_support;

      if ((probe.flags & IDE_PROBE_AVAIL))
        ide_make_device (proc, rctx, probe);
    }
  }

  return true;
}