#include <device/device.h>
#include <device/dos.h>
#include <device/partdrv.h>
#include <device/partitions.h>
#include <device/ramdrv.h>
#include <device/terminal.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 <mm/liballoc.h>
#include <status.h>
#include <sync/spin_lock.h>
#include <sys/debug.h>

#if defined(__x86_64__)
#include <device/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) {
  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);
  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);

  if (found_link == NULL)
    return NULL;

  return hash_entry (found_link, struct device, device_table_link);
}

struct device* device_create (const char* key, device_op_func_t* ops, size_t ops_len,
                              device_init_func_t init, device_fini_func_t fini, void* arg) {
  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));

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

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

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

  spin_lock (&device_table.lock);

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

  spin_unlock (&device_table.lock);

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

  return device;
}

void terminal_device_init (void) {
  device_op_func_t ops[] = {
      [TERMINAL_PUTSTR] = &terminal_putstr,
  };
  device_create ("TERMINAL", ops, lengthof (ops), &terminal_init, &terminal_fini, NULL);
}

void ramdisk_device_init (void) {
  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,
  };

  struct limine_module_response* module = limine_module_request.response;

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

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

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

  struct ramdrv_init init = {
      .total_size = file->size,
      .sector_size = 512,
      .buffer = file->address,
  };
  device_create ("RD", ops, lengthof (ops), &ramdrv_init, &ramdrv_fini, &init);
}

void vdisk_device_init (void) {
  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,
  };

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

  struct dos_mbr mbr;
  memset (&mbr, 0, sizeof (mbr));

  mbr.valid_sign[0] = 0x55;
  mbr.valid_sign[1] = 0xAA;
  mbr.ptes[0].start_lba = 0;
  mbr.ptes[0].sector_count = div_align_up (init.total_size, init.sector_size);

  struct device* vdisk =
      device_create ("VD", ops, lengthof (ops), &ramdrv_init, &ramdrv_fini, &init);

  size_t sector = 0;
  size_t sector_count = 1;

  spin_lock (&vdisk->lock);
  device_op (vdisk, XDRV_WRITE, NULL, NULL, &sector, &sector_count, &mbr);
  spin_unlock (&vdisk->lock);
}

#if defined(__x86_64__)
void ps2kb_device_init (void) {
  device_op_func_t ops[] = {
      [KB_READ_KEY] = &ps2kb_read_key,
  };
  device_create ("KB", ops, lengthof (ops), &ps2kb_init, &ps2kb_fini, NULL);
}
#endif

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

  terminal_device_init ();
  ramdisk_device_init ();
  vdisk_device_init ();

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