#include <stdlib.h>
#include <stdbool.h>
#include <assert.h>
#include <string.h>
#include "tree.h"


struct tree_node* tree_make(const struct tree_pair *pairs, size_t npairs, size_t disk_size) {
  size_t datalen = 0;
  for (size_t i = 0; i < npairs; i++) {
    datalen += strlen(pairs[i].name) + 1 + sizeof(struct tree_node*);
  }
  datalen += 1;

  struct tree_node *const node = malloc(sizeof(struct tree_node));
  node->data = malloc(datalen);

  size_t total_files = 0;
  size_t total_dirs = 0;
  size_t total_size = disk_size;

  size_t cursor = 0;
  for (size_t i = 0; i < npairs; i++) {
    if (pairs[i].sub == NULL) {
      total_files++;
    } else {
      total_files += pairs[i].sub->total_files;
      total_dirs += 1 + pairs[i].sub->total_dirs;  // include this directory too
      total_size += pairs[i].sub->total_size;
    }

    uint8_t *dst = node->data + cursor;
    cursor += (uint8_t*)stpcpy((char*)dst, pairs[i].name) - dst;
    cursor++;
    *(struct tree_node**)(node->data + cursor) = pairs[i].sub;
    cursor += sizeof(struct tree_node*);
  }
  node->data[cursor] = 0;

  node->total_files = total_files;
  node->total_dirs = total_dirs;
  node->total_size = total_size;

  return node;
}

#define LOOP_NODE(node_, namevar_, subvar_, ...) do { \
    size_t _cursor = 0; \
    while (true) { \
      namevar_ = (char*)((node_)->data + _cursor); \
      size_t _len = strlen((char*)((node_)->data + _cursor)); \
      if (_len == 0) break; \
      _cursor += _len + 1; \
      subvar_ = *(struct tree_node**)((node_)->data + _cursor); \
      _cursor += sizeof(struct tree_node*); \
      { __VA_ARGS__ }; \
    } \
  } while (0)

void tree_free(struct tree_node *node) {
  char *name; (void)name;
  struct tree_node *sub;
  LOOP_NODE(node, name, sub, {
    if (sub != NULL) tree_free(sub);
  });
  free(node->data);
  free(node);
}

size_t tree_num_children(const struct tree_node *node) {
  char *name; (void)name;
  struct tree_node *sub; (void)sub;
  size_t count = 0;
  LOOP_NODE(node, name, sub, {
    count++;
  });
  return count;
}

struct tree_pair tree_index(struct tree_node *node, size_t index) {
  char *name;
  struct tree_node *sub;
  size_t count = 0;
  LOOP_NODE(node, name, sub, {
    if (count == index) return (struct tree_pair){.name = name, .sub = sub};
    count++;
  });
  assert(false);
}

struct tree_iter tree_iter_start(struct tree_node *node) {
  return (struct tree_iter){.node = node, .off = 0};
}

bool tree_iter_next(struct tree_iter *iter, struct tree_pair *dst) {
  struct tree_node *node = iter->node;
  if (node->data[iter->off] == 0) return false;

  dst->name = (char*)(node->data + iter->off);
  const size_t namelen = strlen(dst->name);
  dst->sub = *(struct tree_node**)(node->data + iter->off + namelen + 1);
  iter->off += namelen + 1 + sizeof(struct tree_node*);
  return true;
}