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

DEFINE_DEVICE_INIT (partdrv_init) {
  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;
}

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

  free (partdrv);
}

DEFINE_DEVICE_OP (partdrv_read) {
  uint64_t fs;

  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, &fs);
  int ret = device_op (super, XDRV_READ, proc, rctx, &fs, &sector, &sector_count, buffer);
  spin_unlock (&super->lock, fs);

  return ret;
}

DEFINE_DEVICE_OP (partdrv_write) {
  uint64_t fs;

  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, &fs);
  int ret = device_op (super, XDRV_WRITE, proc, rctx, &fs, &sector, &sector_count, buffer);
  spin_unlock (&super->lock, fs);

  return ret;
}

DEFINE_DEVICE_OP (partdrv_get_device_type) {
  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  int* device_type = (int*)a1;

  *device_type = XDRV_TYPE_PARTDRV;

  return ST_OK;
}

DEFINE_DEVICE_OP (partdrv_get_sector_size) {
  uint64_t fs;

  if (a1 == NULL)
    return -ST_BAD_ADDRESS_SPACE;

  size_t* secsize = (size_t*)a1;

  struct partdrv* partdrv = device->udata;

  spin_lock (&partdrv->super->lock, &fs);
  device_op (partdrv->super, XDRV_GET_SECTOR_SIZE, proc, rctx, &fs, secsize);
  spin_unlock (&partdrv->super->lock, fs);

  return ST_OK;
}

DEFINE_DEVICE_OP (partdrv_get_size) {
  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;
}