#include <device/device.h>
#include <device/storage/partdrv.h>
#include <device/storage/partitions.h>
#include <device/storage/ramdrv.h>
#include <device/sys/debugconsole.h>
#include <device/sys/terminal.h>
#include <device_info.h>
#include <devices.h>
#include <id/id_alloc.h>
#include <libk/align.h>
#include <libk/fieldlengthof.h>
#include <libk/hash.h>
#include <libk/rbtree.h>
#include <libk/std.h>
#include <libk/string.h>
#include <limine/requests.h>
#include <lz4/lz4.h>
#include <lz4/lz4frame.h>
#include <mm/malloc.h>
#include <status.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>

#if defined(__x86_64__)
#include <device/pci/pci.h>
#include <device/ps2/ps2_kb.h>
#endif

struct device_table {
  struct hash_node_link* device_buckets[1024];
  spin_lock_t lock;
};

static struct device_table device_table;

struct device* device_find(const char* key) {
  uint64_t fdt;

  struct hash_node_link* found_link = NULL;
  size_t key_len = strlen_null(key);
  uint32_t hash = hash_fnv32(key, key_len);

  spin_lock(&device_table.lock, &fdt);

  hash_find(&device_table, key, key_len, hash, lengthof(device_table.device_buckets),
            device_buckets, struct device, device_table_link, key, found_link);

  spin_unlock(&device_table.lock, fdt);

  if (found_link == NULL) {
    return NULL;
  }

  struct device* device = hash_entry(found_link, struct device, device_table_link);

  return device;
}

struct device* device_create(int type, const char* key, device_op_func_t* ops, size_t ops_len,
                             device_init_func_t init, device_fini_func_t fini, void* arg,
                             struct proc* proc, struct reschedule_ctx* rctx) {
  uint64_t fdt;

  if (ops_len >= fieldlengthof(struct device, ops))
    return NULL;

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

  if (device == NULL)
    return NULL;

  memset(device, 0, sizeof(*device));

  device->init = init;
  device->fini = fini;
  memcpy(device->key, key, strlen_null(key));
  device->type = type;

  if (ops != NULL)
    memcpy(device->ops, ops, ops_len * sizeof(device_op_func_t));

  if (!device->init(device, arg, proc, rctx)) {
    free(device);
    return NULL;
  }

  uint32_t device_hash = hash_fnv32(device->key, strlen_null(device->key));

  spin_lock(&device_table.lock, &fdt);

  hash_insert(&device_table, &device->device_table_link, device_hash,
              lengthof(device_table.device_buckets), device_buckets);

  spin_unlock(&device_table.lock, fdt);

  DEBUG("Created device %s\n", device->key);

  return device;
}

void device_delete(const char* key, struct proc* proc, struct reschedule_ctx* rctx) {
  uint64_t fdt, fsd, fd;

  size_t key_len = strlen_null(key);
  uint32_t hash = hash_fnv32(key, key_len);
  struct hash_node_link* found_link;

  spin_lock(&device_table.lock, &fdt);

  hash_find(&device_table, key, key_len, hash, lengthof(device_table.device_buckets),
            device_buckets, struct device, device_table_link, key, found_link);

  if (found_link == NULL) {
    spin_unlock(&device_table.lock, fdt);
    return;
  }

  struct device* device = hash_entry(found_link, struct device, device_table_link);

  if (device == NULL) {
    spin_unlock(&device_table.lock, fdt);
    return;
  }

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

  hash_delete(&device_table, key, strlen_null(key), hash, lengthof(device_table.device_buckets),
              device_buckets, struct device, device_table_link, key, found_link);

  struct list_node_link *subdevice_link, *tmp_subdevice_link;
  list_foreach(device->subdevices, subdevice_link, tmp_subdevice_link) {
    struct device* subdevice = list_entry(subdevice_link, struct device, subdevices_link);

    spin_lock(&subdevice->lock, &fsd);

    list_remove(device->subdevices, &subdevice->subdevices_link);

    device->fini(subdevice, proc, rctx);

    spin_unlock(&subdevice->lock, fsd);

    free(subdevice);
  }

  device->fini(device, proc, rctx);

  spin_unlock(&device->lock, fd);

  free(device);

  spin_unlock(&device_table.lock, fdt);
}

size_t device_populate_device_infos(struct device_info* infos, size_t count) {
  uint64_t fdt, fd;

  if (count >= DEVICES_MAX)
    count = DEVICES_MAX;

  size_t j = 0;

  spin_lock(&device_table.lock, &fdt);

  for (size_t i = 0; i < lengthof(device_table.device_buckets); i++) {
    struct hash_node_link* hash_link = device_table.device_buckets[i];

    while (hash_link != NULL && j < count) {
      struct device* device = hash_entry(hash_link, struct device, device_table_link);
      struct device_info* info = &infos[j];

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

      memcpy(info->key, device->key, sizeof(info->key));
      info->type = device->type;

      spin_unlock(&device->lock, fd);

      hash_link = hash_link->next;
      j++;
    }
  }

  spin_unlock(&device_table.lock, fdt);

  return j;
}

static void debugconsole_device_init(void) {
  struct reschedule_ctx rctx;
  memset(&rctx, 0, sizeof(rctx));

  static device_op_func_t ops[] = {
      [DEBUGCONSOLE_PUTSTR] = &debugconsole_putstr,
  };

  device_create(DEVICE_TYPE_DEBUGCONSOLE, "debugconsole", ops, lengthof(ops), &debugconsole_init,
                &debugconsole_fini, NULL, thiscpu->kproc, &rctx);
}

static void terminal_device_init(void) {
  struct reschedule_ctx rctx;
  memset(&rctx, 0, sizeof(rctx));

  static device_op_func_t ops[] = {
      [TERMINAL_PUTSTR] = &terminal_putstr,
      [TERMINAL_DIMENSIONS] = &terminal_dimensions,
  };

  device_create(DEVICE_TYPE_TERMINAL, "terminal", ops, lengthof(ops), &terminal_init,
                &terminal_fini, NULL, thiscpu->kproc, &rctx);
}

static void sys_device_init(void) {
  struct reschedule_ctx rctx;
  memset(&rctx, 0, sizeof(rctx));

  static device_op_func_t ops[] = {
      [XDRV_GET_SIZE] = &ramdrv_get_size,
      [XDRV_GET_SECTOR_SIZE] = &ramdrv_get_sector_size,
      [XDRV_GET_DEVICE_TYPE] = &ramdrv_get_device_type,
      [XDRV_READ] = &ramdrv_read,
      [XDRV_WRITE] = &ramdrv_write,
      [XDRV_PARTITION_RESCAN] = &ramdrv_partition_rescan,
  };

  struct limine_module_response* module = limine_module_request.response;

  const char* dist_path = "/boot/mop3dist.tar.lz4";
  struct limine_file* file = NULL;

  for (size_t i = 0; i < module->module_count; i++) {
    file = module->modules[i];

    if (strncmp(file->path, dist_path, strlen(dist_path)) == 0)
      break;
  }

  const size_t unpack_buffer_size = 1024 * 1024 * 16;
  uint8_t* unpack_buffer = malloc(unpack_buffer_size);

  LZ4F_decompressionContext_t dctx;

  LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION);

  size_t consumed = file->size;
  size_t produced = unpack_buffer_size;

  LZ4F_decompress(dctx, unpack_buffer, &produced, file->address, &consumed, NULL);

  struct ramdrv_init init = {
      .total_size = produced,
      .sector_size = 512,
      .buffer = unpack_buffer,
  };

  device_create(DEVICE_TYPE_DRIVE, "sys0", ops, lengthof(ops), &ramdrv_init, &ramdrv_fini, &init,
                thiscpu->kproc, &rctx);

  LZ4F_freeDecompressionContext(dctx);
}

static void temp_device_init(void) {
  struct reschedule_ctx rctx;
  memset(&rctx, 0, sizeof(rctx));

  static device_op_func_t ops[] = {
      [XDRV_GET_SIZE] = &ramdrv_get_size,
      [XDRV_GET_SECTOR_SIZE] = &ramdrv_get_sector_size,
      [XDRV_GET_DEVICE_TYPE] = &ramdrv_get_device_type,
      [XDRV_READ] = &ramdrv_read,
      [XDRV_WRITE] = &ramdrv_write,
      [XDRV_PARTITION_RESCAN] = &ramdrv_partition_rescan,
  };

  struct ramdrv_init init = {
      .total_size = 1024 * 1024 * 20,
      .sector_size = 512,
  };

  device_create(DEVICE_TYPE_DRIVE, "temp0", ops, lengthof(ops), &ramdrv_init, &ramdrv_fini, &init,
                thiscpu->kproc, &rctx);
}

#if defined(__x86_64__)
static void ps2kb_device_init(void) {
  struct reschedule_ctx rctx;
  memset(&rctx, 0, sizeof(rctx));

  device_op_func_t ops[] = {
      [KB_READ_KEY] = &ps2kb_read_key,
  };

  device_create(DEVICE_TYPE_KEYBOARD, "kb", ops, lengthof(ops), &ps2kb_init, &ps2kb_fini, NULL,
                thiscpu->kproc, &rctx);
}
#endif

void devices_init(void) {
  memset(&device_table, 0, sizeof(device_table));

  terminal_device_init();
  sys_device_init();
  temp_device_init();
  debugconsole_device_init();

#if defined(__x86_64__)
  ps2kb_device_init();
  pci_init();
#endif
}