#include <initrd.h>
#include <utils.h>
#include <stddef.h>
#include <dtb.h>
#include <kmalloc.h>
#include <random.h>
#include <string.h>
#include <uart.h>

#define nullnode ((file_node_t *)0)

fdt_callback_t initrd_dtb_cb = {
  .name = "chosen",
  .func = initrd_fdt_callback,
};

void *initrd_start = 0x0;
void *initrd_end   = 0x0;

void initrd_fdt_callback(const vector_t *props)
{
  for (int i = 0; i < (int)props->size; ++i) {
    if (!strcmp(VEC_AT(fdt_prop_t, props, i)->name, "linux,initrd-start"))
      initrd_start = (void *)(uint64_t)ntoh32(
        *(uint32_t *)VEC_AT(fdt_prop_t, props, i)->value);
    if (!strcmp(VEC_AT(fdt_prop_t, props, i)->name, "linux,initrd-end"))
      initrd_end = (void *)(uint64_t)ntoh32(
        *(uint32_t *)VEC_AT(fdt_prop_t, props, i)->value);
  }
}

void _init_node(file_node_t *node) {
  node->l = nullnode;
  node->r = nullnode;
  node->rand = random();
  node->node_size = 1;
}

void _pull_from(file_node_t *to, file_node_t *from)
{
  if (!from)
    return;

  to->node_size += from->node_size;
}

void _pull(file_node_t *node)
{
  node->node_size = 1;
  _pull_from(node, node->l);
  _pull_from(node, node->r);
}

file_node_t *_merge(file_node_t *a, file_node_t *b)
{
  if (!a || !b)
    return a ?: b;

  if (a->rand < b->rand) {
    a->r = _merge(a->r, b);
    _pull(a);
    return a;
  }
  b->l = _merge(a, b->l);
  _pull(b);
  return b;
}

void _split(file_node_t *rt, const char *s, file_node_t **a, file_node_t **b)
{
  if (!rt) {
    *a = *b = nullnode;
    return;
  }

  if (strcmp(rt->filename, s) < 0) {
    *a = rt;
    _split((*a)->r, s, &(*a)->r, b);
    _pull(*a);
  } else {
    *b = rt;
    _split((*b)->l, s, a, &(*b)->l);
    _pull(*b);
  }
}

cpio_newc_header_t *initrd_cur;

file_node_t *initrd_init(void)
{
  cpio_newc_header_t *header;
  char *filename;
  uint8_t *filecontent;

  initrd_cur = (void *)initrd_start;

  file_node_t *root = nullnode, *ltr, *rtr;
  for (;;) {
    header = BUMP(cpio_newc_header_t, cpio_newc_header_t, initrd_cur);

    filename = (void *)initrd_cur;
    if (!strcmp(filename, "TRAILER!!!"))
      break;
    initrd_cur = ALIGN4((uint64_t)initrd_cur + atoh32(header->c_namesize));

    filecontent = (void *)initrd_cur;
    initrd_cur = ALIGN4((uint64_t)initrd_cur + atoh32(header->c_filesize));

    file_node_t *newnode = kmalloc(sizeof(file_node_t));
    _init_node(newnode);
    newnode->ino       = atoh32(header->c_ino);
    newnode->mode      = atoh32(header->c_mode);
    newnode->uid       = atoh32(header->c_uid);
    newnode->gid       = atoh32(header->c_gid);
    newnode->nlink     = atoh32(header->c_nlink);
    newnode->mtime     = atoh32(header->c_mtime);
    newnode->filesize  = atoh32(header->c_filesize);
    newnode->devmajor  = atoh32(header->c_devmajor);
    newnode->devminor  = atoh32(header->c_devminor);
    newnode->rdevmajor = atoh32(header->c_rdevmajor);
    newnode->rdevminor = atoh32(header->c_rdevminor);
    newnode->namesize  = atoh32(header->c_namesize);

    newnode->filename = filename;
    newnode->filecontent = filecontent;

    DEBUG_s(filename);

    _split(root, filename, &ltr, &rtr);
    root = _merge(ltr, _merge(newnode, rtr));
  }
  
  return root;
}

file_node_t *_node_bs(file_node_t *cur, const char *s)
{
  if (!cur) 
    return nullnode;

  int cmp = strcmp(cur->filename, s);
  if (cmp < 0)
    return _node_bs(cur->r, s);
  if (cmp > 0)
    return _node_bs(cur->l, s);
  return cur; // cmp == 0
}

file_node_t *initrd_get(file_node_t *tr, const char *filename)
{
  return _node_bs(tr, filename);
}

void initrd_traverse(file_node_t *tr, initrd_callback_func_t func)
{
  if (!tr)
    return;

  initrd_traverse(tr->l, func);
  func(tr);
  initrd_traverse(tr->r, func);
}