#include <device/device.h>
#include <device/storage/partdrv.h>
#include <device/storage/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 <proc/proc.h>
#include <proc/reschedule.h>
#include <status.h>
#include <sys/debug.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,
      [XDRV_PARTITION_RESCAN] = &partdrv_partition_rescan,
  };

  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) {
    return ret;
  }

  if (!(mbr.valid_sign[0] == 0x55 && mbr.valid_sign[1] == 0xAA)) {
    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);

    struct device* part_dev = device_create(DEVICE_TYPE_DRIVE, key, ops, lengthof(ops),
                                            &partdrv_init, &partdrv_fini, &init, proc, rctx);

    list_append(device->subdevices, &part_dev->subdevices_link);
  }

  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);
}