#include <amd64/io.h>
#include <device/device.h>
#include <device/idedrv.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 <sys/debug.h>

static atomic_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 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", atomic_fetch_add (&ide_counter, 1));

  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,
  };
  device_create (device_key, ops, lengthof (ops), &idedrv_init, &idedrv_fini, &init);
}

bool pci_ide_init (struct pci_info pci_info) {
  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)) {
    pcmd = (uint16_t)(pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR0) & 0xFFFC);
    pctrl = (uint16_t)(pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR1) & 0xFFFC);

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

  if ((progif & 0x04)) {
    scmd = (uint16_t)(pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR2) & 0xFFFC);
    sctrl = (uint16_t)(pci_read32 (pci_info.bus, pci_info.slot, pci_info.func, PCI_BAR3) & 0xFFFC);

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

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

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

    if (cmd == 0)
      continue;

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

      if ((probe.flags & IDE_PROBE_AVAIL))
        ide_make_device (probe);
    }
  }

  return true;
}