#include <device/device.h>
#include <device/partitions.h>
#include <device/ramdrv.h>
#include <devices.h>
#include <libk/align.h>
#include <libk/std.h>
#include <libk/string.h>
#include <mm/liballoc.h>
#include <status.h>
#include <sys/debug.h>

bool ramdrv_init (struct device* device, void* arg) {
  struct ramdrv_init* init = arg;

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

  if (ramdrv == NULL)
    return false;

  ramdrv->sector_size = init->sector_size;
  ramdrv->total_size = align_up (init->total_size, ramdrv->sector_size);

  ramdrv->buffer = malloc (ramdrv->total_size);

  if (ramdrv->buffer == NULL) {
    free (ramdrv);
    return false;
  }

  memset (ramdrv->buffer, 0, ramdrv->total_size);

  if (init->buffer != NULL)
    memcpy (ramdrv->buffer, init->buffer, init->total_size);

  device->udata = ramdrv;

  return true;
}

void ramdrv_fini (struct device* device) {
  struct ramdrv* ramdrv = device->udata;

  free (ramdrv->buffer);
  free (ramdrv);
}

int ramdrv_get_device_type (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                            void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)device, (void)a2, (void)a3, (void)a4;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  int* device_type = (int*)a1;

  *device_type = XDRV_TYPE_RAMDRV;

  return ST_OK;
}

int ramdrv_get_size (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                     void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a2, (void)a3, (void)a4;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* size = (size_t*)a1;

  struct ramdrv* ramdrv = device->udata;

  *size = ramdrv->total_size;

  return ST_OK;
}

int ramdrv_get_sector_size (struct device* device, struct proc* proc, struct reschedule_ctx* rctx,
                            void* a1, void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a2, (void)a3, (void)a4;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* secsize = (size_t*)a1;

  struct ramdrv* ramdrv = device->udata;

  *secsize = ramdrv->sector_size;

  return ST_OK;
}

int ramdrv_read (struct device* device, struct proc* proc, struct reschedule_ctx* rctx, void* a1,
                 void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a4;

  if (a1 == NULL || a2 == NULL || a3 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t sector = *(size_t*)a1;
  size_t sector_count = *(size_t*)a2;
  uint8_t* buffer = a3;

  struct ramdrv* ramdrv = device->udata;
  size_t pos = sector * ramdrv->sector_size;
  size_t size = sector_count * ramdrv->sector_size;

  memcpy (buffer, (void*)(((uintptr_t)ramdrv->buffer) + pos), size);

  return ST_OK;
}

int ramdrv_write (struct device* device, struct proc* proc, struct reschedule_ctx* rctx, void* a1,
                  void* a2, void* a3, void* a4) {
  (void)proc, (void)rctx, (void)a4;

  if (a1 == NULL || a2 == NULL || a3 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t sector = *(size_t*)a1;
  size_t sector_count = *(size_t*)a2;
  uint8_t* buffer = a3;

  struct ramdrv* ramdrv = device->udata;
  size_t pos = sector * ramdrv->sector_size;
  size_t size = sector_count * ramdrv->sector_size;

  memcpy ((void*)(((uintptr_t)ramdrv->buffer) + pos), buffer, size);

  return ST_OK;
}
