#include <stdbool.h>
#include <stdint.h>
#include <stddef.h>
#include "storedev/storedev.h"
#include "diskpart/mbr.h"
#include "diskpart/diskpart.h"
#include "compiler/attr.h"
#include "util/util.h"
#include "dev/dev.h"
#include "std/string.h"
#include "errors.h"
#include "kprintf.h"

typedef struct {
  uint8_t driveattrs;
  uint8_t chs_start_addr[3];
  uint8_t parttype;
  uint8_t chs_last_sect_addr[3];
  uint32_t startlba;
  uint32_t sectorcount;
} PACKED PTE;

typedef struct {
  uint8_t bootcode[440];
  uint8_t signature[4];
  uint8_t resv[2];
  PTE pte1;
  PTE pte2;
  PTE pte3;
  PTE pte4;
  uint8_t validsign[2];
} PACKED MBR;

void diskpart_mbr_expose_partitions(StoreDev *sd, MBR *mbr) {
  PTE *parts[4] = { &mbr->pte1, &mbr->pte2, &mbr->pte3, &mbr->pte4 };

  for (size_t i = 0; i < 4; i++) {
    char *name = NULL;
    spinlock_acquire(&DEVTABLE.spinlock);
    for (size_t i = 0; i < LEN(DEVTABLE.devs); i++) {
      Dev *dev = &DEVTABLE.devs[i];
      if (dev->extra == sd) {
        name = (char *)dev->ident;
      }
    }
    spinlock_release(&DEVTABLE.spinlock);

    PartSdInitExtra extra = {
      .parent = sd,
      .sectorcount = parts[i]->sectorcount,
      .startsector = parts[i]->startlba,
    };

    char name1[100];
    ksprintf(name1, "%sp%zu", name, i+1);
    kprintf("%s\n", name1);
    storedev_create(STOREDEV_PARTSD, name1, &extra);
  }
}

bool diskpart_mbr_probe(StoreDev *sd) {
  LOG("diskpart", "probing for MBR\n");

  MBR mbr;
  memset(&mbr, 0, sizeof(mbr));
  sd->read(sd, (uint8_t *const)&mbr, 0, 0, sizeof(mbr));

  if (!(mbr.validsign[0] == 0x55 && mbr.validsign[1] == 0xAA)) {
    return false;
  }

  LOG("diskpart", "got valid MBR on %s\n", storedev_strings[sd->sdtype]);

  diskpart_mbr_expose_partitions(sd, &mbr);

  return true;
}