#include <device/device.h>
#include <device/partdrv.h>
#include <devices.h>
#include <libk/std.h>
#include <mm/liballoc.h>
#include <proc/proc.h>
#include <proc/reschedule.h>
#include <status.h>

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

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

  if (partdrv == NULL)
    return false;

  partdrv->start_sector = init->start_sector;
  partdrv->super = init->super;
  partdrv->total_size = init->total_size;

  device->udata = partdrv;

  return true;
}

void partdrv_fini (struct device* device) {
  struct partdrv* partdrv = device->udata;

  free (partdrv);
}

int partdrv_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;

  struct partdrv* partdrv = device->udata;
  struct device* super = partdrv->super;

  size_t sector = *(size_t*)a1 + partdrv->start_sector;
  size_t sector_count = *(size_t*)a2;
  uint8_t* buffer = a3;

  spin_lock (&super->lock);
  int ret = device_op (super, XDRV_READ, NULL, NULL, &sector, &sector_count, buffer);
  spin_unlock (&super->lock);

  return ret;
}

int partdrv_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;

  struct partdrv* partdrv = device->udata;
  struct device* super = partdrv->super;

  size_t sector = *(size_t*)a1 + partdrv->start_sector;
  size_t sector_count = *(size_t*)a2;
  uint8_t* buffer = a3;

  spin_lock (&super->lock);
  int ret = device_op (super, XDRV_WRITE, NULL, NULL, &sector, &sector_count, buffer);
  spin_unlock (&super->lock);

  return ret;
}

int partdrv_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_PARTDRV;

  return ST_OK;
}

int partdrv_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 partdrv* partdrv = device->udata;

  spin_lock (&partdrv->super->lock);
  device_op (partdrv->super, XDRV_GET_SECTOR_SIZE, NULL, NULL, secsize);
  spin_unlock (&partdrv->super->lock);

  return ST_OK;
}

int partdrv_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 partdrv* partdrv = device->udata;

  *size = partdrv->total_size;

  return ST_OK;
}