#include <device/device.h>
#include <device/partdrv.h>
#include <device/partitions.h>
#include <devices.h>
#include <libk/fieldsizeof.h>
#include <libk/lengthof.h>
#include <libk/printf.h>
#include <libk/std.h>
#include <libk/string.h>
#include <status.h>
#include <sys/debug.h>
#include <proc/proc.h>
#include <proc/reschedule.h>

static int device_probe_partitions_dos (struct proc* proc, struct reschedule_ctx* rctx, struct device* device) {
  struct dos_mbr mbr;
  memset (&mbr, 0, sizeof (mbr));
  size_t sector = 0;
  size_t sector_count = 1;
  size_t sector_size;

  device_op_func_t ops[] = {
      [XDRV_GET_SIZE] = &partdrv_get_size,
      [XDRV_GET_SECTOR_SIZE] = &partdrv_get_sector_size,
      [XDRV_GET_DEVICE_TYPE] = &partdrv_get_device_type,
      [XDRV_READ] = &partdrv_read,
      [XDRV_WRITE] = &partdrv_write,
  };

  spin_lock (&device->lock);

  device_op (device, XDRV_GET_SECTOR_SIZE, proc, rctx, &sector_size);

  int ret = device_op (device, XDRV_READ, proc, rctx, &sector, &sector_count, &mbr);

  if (ret < 0) {
    spin_unlock (&device->lock);
    return ret;
  }

  if (!(mbr.valid_sign[0] == 0x55 && mbr.valid_sign[1] == 0xAA)) {
    spin_unlock (&device->lock);
    return -ST_PARTITION_ERROR;
  }

  for (size_t i = 0; i < lengthof (mbr.ptes); i++) {
    struct dos_pte* pte = &mbr.ptes[i];

    struct partdrv_init init = {
        .start_sector = pte->start_lba,
        .total_size = pte->sector_count * sector_size,
        .super = device,
    };

    char key[fieldsizeof (struct device, key)];
    memset (key, 0, sizeof (key));
    snprintf (key, sizeof (key), "%sp%zu", device->key, i);

    device_create (key, ops, lengthof (ops), &partdrv_init, &partdrv_fini, &init, proc, rctx);
  }

  spin_unlock (&device->lock);

  return ST_OK;
}

int device_probe_partitions (struct proc* proc, struct reschedule_ctx* rctx, struct device* device) { return device_probe_partitions_dos (proc, rctx, device); }