#include "src/kernel/nvme.h"
#include "src/kernel/log.h"
#include "src/kernel/panic.h"
#include "src/kernel/pci.h"
#include "src/lib/layout.h"
#include "src/lib/memory.h"
#include "src/lib/util.h"

#define SECTOR_SIZE 512

#define NVME_PCI_CLASS 0x01
#define NVME_PCI_SUBCLASS 0x08

typedef volatile struct __attribute__((packed)) {
  uint64_t cap;
  uint32_t vs;
  uint32_t intms;
  uint32_t intmc;
  uint32_t cc;
  uint32_t rsvd;
  uint32_t csts;
  uint32_t nssr;
  uint32_t aqa;
  uint64_t asq;
  uint64_t acq;
} nvme_regs_t;

#define NVME_REGS_NVME_CS_SUPPORTED(r) BITS_R(r->cap, 37, 37)
#define NVME_REGS_MQES(r) BITS_R(r->cap, 15, 0)
#define NVME_REGS_IOCQES_R(r) BITS_R(r->cc, 23, 20)
#define NVME_REGS_IOCQES_W(r, v) BITS_W(r->cc, 23, 20, v)
#define NVME_REGS_IOSQES_R(r) BITS_R(r->cc, 19, 16)
#define NVME_REGS_IOSQES_W(r, v) BITS_W(r->cc, 19, 16, v)
#define NVME_REGS_CSS_R(r) BITS_R(r->cc, 6, 4)
#define NVME_REGS_CSS_W(r, v) BITS_W(r->cc, 6, 4, v)
#define NVME_REGS_EN_R(r) BITS_R(r->cc, 0, 0)
#define NVME_REGS_EN_W(r, v) BITS_W(r->cc, 0, 0, v)
#define NVME_REGS_RDY(r) BITS_R(r->csts, 0, 0)
#define NVME_REGS_ACQS_R(r) BITS_R(r->aqa, 27, 16)
#define NVME_REGS_ACQS_W(r, v) BITS_W(r->aqa, 27, 16, v)
#define NVME_REGS_ASQS_R(r) BITS_R(r->aqa, 11, 0)
#define NVME_REGS_ASQS_W(r, v) BITS_W(r->aqa, 11, 0, v)

static nvme_regs_t *regs = (nvme_regs_t *)KERNEL_VIRTUAL_NVME;
static volatile uint32_t *db_regs = (uint32_t *)(KERNEL_VIRTUAL_NVME + 0x1000); // TODO Get stride from caps.

#define QUEUE_DEPTH 2

typedef volatile struct __attribute__((packed)) {
  uint8_t opc;
  uint8_t flags;
  uint16_t cid;
  uint32_t nsid;
  uint64_t reserved;
  uint64_t mptr;
  uint64_t prp1;
  uint64_t prp2;
  uint32_t cdw10;
  uint32_t cdw11;
  uint32_t cdw12;
  uint32_t cdw13;
  uint32_t cdw14;
  uint32_t cdw15;
} nvme_sqe_t;

typedef volatile struct __attribute__((packed)) {
  uint32_t dw0;
  uint32_t reserved;
  uint16_t sqhd;
  uint16_t sqid;
  uint16_t cid;
  uint16_t status;
} nvme_cqe_t;

#define NVME_CQE_STATUS(e) BITS_R(e->status, 15, 1)
#define NVME_CQE_P(e) BITS_R(e->status, 0, 0)

static uint8_t admin_cq_phase = 1;
static nvme_sqe_t admin_sq[QUEUE_DEPTH] __attribute__((aligned(PAGE_SIZE)));
static uint16_t admin_sq_tail = 0;
static nvme_cqe_t admin_cq[QUEUE_DEPTH] __attribute__((aligned(PAGE_SIZE)));
static uint16_t admin_cq_head = 0;
static uint8_t io_cq_phase = 1;
static nvme_sqe_t io_sq[QUEUE_DEPTH] __attribute__((aligned(PAGE_SIZE)));
static uint16_t io_sq_tail = 0;
static nvme_cqe_t io_cq[QUEUE_DEPTH] __attribute__((aligned(PAGE_SIZE)));
static uint16_t io_cq_head = 0;

static void admin_exec_sync(nvme_sqe_t *cmd) {
  cmd->cid = admin_sq_tail;

  LOG_LN_STEP("Queueing submission...");
  LOG_VAL_TRACE(cmd->opc, "%hhx");
  LOG_VAL_TRACE(cmd->nsid, "%x");
  LOG_VAL_TRACE(cmd->cdw10, "%x");
  LOG_VAL_TRACE(cmd->cdw11, "%x");
  LOG_VAL_TRACE(cmd->cdw12, "%x");
  LOG_VAL_TRACE(cmd->prp1, "%lx");
  LOG_VAL_TRACE(cmd->prp2, "%lx");

  memory_copy(cmd, sizeof(nvme_sqe_t), &admin_sq[admin_sq_tail]);

  LOG_LN_STEP("Advancing submission doorbell...");
  db_regs[0] = admin_sq_tail = (admin_sq_tail + 1) % QUEUE_DEPTH;
  LOG_VAL_TRACE(admin_sq_tail, "%u");

  LOG_LN_STEP("Waiting for completion...");
  nvme_cqe_t *cmp = &admin_cq[admin_cq_head];
  while (NVME_CQE_P(cmp) != admin_cq_phase)
    ;

  LOG_LN_STEP("Received completion...");
  LOG_VAL_TRACE(cmp->dw0, "%x");
  LOG_VAL_TRACE(cmp->reserved, "%x");
  LOG_VAL_TRACE(cmp->sqhd, "%hx");
  LOG_VAL_TRACE(cmp->sqid, "%hx");
  LOG_VAL_TRACE(cmp->cid, "%hx");
  LOG_VAL_TRACE(cmp->status, "%hx");
  ASSERT(NVME_CQE_STATUS(cmp) == 0, "admin_exec_sync: NVMe command failed.");

  LOG_LN_STEP("Advancing completion doorbell...");
  db_regs[1] = admin_cq_head = (admin_cq_head + 1) % QUEUE_DEPTH;
  admin_cq_phase ^= (admin_cq_head == 0);

  LOG_LN_STEP("Done.");
}

static void io_exec_sync(nvme_sqe_t *cmd) {
  cmd->cid = io_sq_tail;

  LOG_LN_STEP("Queueing submission...");
  LOG_VAL_TRACE(cmd->opc, "%hhx");
  LOG_VAL_TRACE(cmd->nsid, "%x");
  LOG_VAL_TRACE(cmd->cdw10, "%x");
  LOG_VAL_TRACE(cmd->cdw11, "%x");
  LOG_VAL_TRACE(cmd->cdw12, "%x");
  LOG_VAL_TRACE(cmd->prp1, "%x");
  LOG_VAL_TRACE(cmd->prp2, "%x");

  memory_copy(cmd, sizeof(nvme_sqe_t), &io_sq[io_sq_tail]);

  LOG_LN_STEP("Advancing submission doorbell...");
  db_regs[2] = io_sq_tail = (io_sq_tail + 1) % QUEUE_DEPTH;

  LOG_LN_STEP("Waiting for completion...");
  nvme_cqe_t *cmp = &io_cq[io_cq_head];
  while (NVME_CQE_P(cmp) != io_cq_phase)
    ;

  LOG_LN_STEP("Received completion...");
  LOG_VAL_TRACE(cmp->dw0, "%x");
  LOG_VAL_TRACE(cmp->reserved, "%x");
  LOG_VAL_TRACE(cmp->sqhd, "%hx");
  LOG_VAL_TRACE(cmp->sqid, "%hx");
  LOG_VAL_TRACE(cmp->cid, "%hx");
  LOG_VAL_TRACE(cmp->status, "%hx");

  ASSERT(NVME_CQE_STATUS(cmp) == 0, "io_exec_sync: NVMe command failed.");

  LOG_LN_STEP("Advancing completion doorbell...");
  db_regs[3] = io_cq_head = (io_cq_head + 1) % QUEUE_DEPTH;
  io_cq_phase ^= (io_cq_head == 0);

  LOG_LN_STEP("Done.");
}

void nvme_init() {
  LOG_LN_INFO("Searching for suitable device...");
  pci_bdf_t bdf = pci_find_by_class(NVME_PCI_CLASS, NVME_PCI_SUBCLASS, NUL);
  ASSERT(bdf, "nvme_init: No suitable devices found.");

  LOG_LN_INFO("Mapping device to virtual memory...");
  pci_map(bdf, regs, 4);
  ASSERT(NVME_REGS_NVME_CS_SUPPORTED(regs), "nvme_init: NVM command set not supported.");
  ASSERT(QUEUE_DEPTH <= NVME_REGS_MQES(regs), "nvme_init: Command queues are too big.")

  LOG_LN_INFO("Disabling device...");
  regs->cc = 0;
  LOG_VAL_DATA(regs->cc, "%lx");
  while (NVME_REGS_RDY(regs))
    ;
  LOG_VAL_DATA(regs->cc, "%lx");

  LOG_LN_INFO("Configuring device...");
  NVME_REGS_ASQS_W(regs, QUEUE_DEPTH - 1);
  regs->asq = (uint64_t)(VIRT_TO_PHYS(admin_sq));
  NVME_REGS_ACQS_W(regs, QUEUE_DEPTH - 1);
  regs->acq = (uint64_t)(VIRT_TO_PHYS(admin_cq));
  LOG_VAL_DATA(regs->asq, "%lx");
  LOG_VAL_DATA(regs->acq, "%lx");
  LOG_VAL_DATA(regs->aqa, "%lx");

  LOG_LN_INFO("Enabling device...");
  NVME_REGS_IOCQES_W(regs, 4);
  NVME_REGS_IOSQES_W(regs, 6);
  NVME_REGS_CSS_W(regs, 0);
  NVME_REGS_EN_W(regs, 1);
  LOG_VAL_DATA(regs->cc, "%lx");
  while (!NVME_REGS_RDY(regs))
    ;
  LOG_VAL_DATA(regs->cc, "%lx");

  LOG_LN_INFO("Creating I/O completion queue...");
  nvme_sqe_t create_io_cq = {
      .opc = 0x05,
      .prp1 = (uint64_t)VIRT_TO_PHYS(io_cq),
      .cdw10 = ((QUEUE_DEPTH - 1) << 16) | 1,
      .cdw11 = 1,
  };
  admin_exec_sync(&create_io_cq);

  LOG_LN_INFO("Creating I/O submission queue...");
  nvme_sqe_t create_io_sq = {
      .opc = 0x01,
      .prp1 = (uint64_t)VIRT_TO_PHYS(io_sq),
      .cdw10 = ((QUEUE_DEPTH - 1) << 16) | 1,
      .cdw11 = (1 << 16) | 1,
  };
  admin_exec_sync(&create_io_sq);

  LOG_LN_INFO("Done.");
}

void nvme_read_sectors(uint32_t index, uint8_t count, void *to) {
  LOG_LN_STEP("Reading %u sectors at %u...", count, index);

  while (count) {
    uint64_t remainder = PAGE_SIZE - (uint64_t)to % PAGE_SIZE;

    uint8_t i_count = (uint8_t)((remainder + PAGE_SIZE) / SECTOR_SIZE);
    i_count = i_count > count ? count : i_count;

    nvme_sqe_t read_sq = {
        .opc = 0x02,
        .nsid = 1,
        .cdw10 = index,
        .cdw11 = 0,
        .cdw12 = i_count - 1,
        .prp1 = (uint64_t)VIRT_TO_PHYS(to),
        .prp2 = i_count * SECTOR_SIZE > remainder ? (uint64_t)VIRT_TO_PHYS(to + remainder) : 0,
    };
    io_exec_sync(&read_sq);

    index += i_count;
    count -= i_count;
    to = (void *)((uint8_t *)to + i_count * SECTOR_SIZE);
  }

  LOG_LN_STEP("Done.");
}

void nvme_write_sectors(uint32_t index, uint8_t count, const void *from) {
  LOG_LN_STEP("Writing %u sectors at %u...", count, index);

  while (count) {
    uint64_t remainder = PAGE_SIZE - (uint64_t)from % PAGE_SIZE;

    uint8_t i_count = (uint8_t)((remainder + PAGE_SIZE) / SECTOR_SIZE);
    i_count = i_count > count ? count : i_count;

    nvme_sqe_t write_sq = {
        .opc = 0x01,
        .nsid = 1,
        .cdw10 = index,
        .cdw11 = 0,
        .cdw12 = i_count - 1,
        .prp1 = (uint64_t)VIRT_TO_PHYS(from),
        .prp2 = i_count * SECTOR_SIZE > remainder ? (uint64_t)VIRT_TO_PHYS(from + remainder) : 0,
    };
    io_exec_sync(&write_sq);

    index += i_count;
    count -= i_count;
    from = (void *)((uint8_t *)from + i_count * SECTOR_SIZE);
  }

  LOG_LN_STEP("Done.");
}