#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <errno.h>
#include <assert.h>

#include "treetrie.h"
#include "fasthash.h"

static inline tt_hash_t hash(uint32_t tag,
			     uint32_t index,
			     tt_node_idx_t a,
			     tt_node_idx_t b)
{
  uint32_t keyblock[4] = { tag,
			   index,
			   a,
			   b };
  assert(sizeof(keyblock) == 4 * sizeof(uint32_t));
  return (tt_hash_t) fasthash32(keyblock, sizeof(keyblock), 0);
}

inline tt_hash_t tt_hash_node(tt_arena_t *a, tt_node_idx_t i) {
  return hash(a->headers[i].inuse.tag,
	      a->headers[i].inuse.index,
	      a->nodes[i].a,
	      a->nodes[i].b);
}

int tt_arena_init(tt_arena_t *a) {
  a->max_probe = 0;
  a->live_count = 0;
  a->table_length = 16;
  a->table = calloc(a->table_length, sizeof(a->table[0]));
  a->headers = calloc(a->table_length, sizeof(a->headers[0]));
  a->nodes = calloc(a->table_length, sizeof(a->nodes[0]));
  a->free_chain = TT_ERROR;

  if (a->table == NULL || a->headers == NULL || a->nodes == NULL) {
    if (a->table != NULL) free(a->table);
    if (a->headers != NULL) free(a->headers);
    if (a->nodes != NULL) free(a->nodes);
    errno = ENOMEM;
    return -1;
  }

  {
    int i;
    for (i = a->table_length - 1; i >= TT_FIRST_VALID_NODE_IDX; i--) {
      a->headers[i].next_free = a->free_chain;
      a->free_chain = i;
    }
  }

  return 0;
}

static int tt_grow(tt_arena_t *a) {
  assert(0);
}

void tt_arena_done(tt_arena_t *a) {
  free(a->table);
  free(a->headers);
  free(a->nodes);
  memset(a, 0, sizeof(*a));
}

static void recycle_node(tt_arena_t *a, tt_node_idx_t ni) {
  tt_hash_t h;
  int i;

  printf("++++++++++++++++++++++++++++++++++++++++ recycling %d\n", ni);

  assert(ni >= TT_FIRST_VALID_NODE_IDX);
  h = tt_hash_node(a, ni);

  if (a->headers[ni].inuse.tag == TT_TAG_LEAF) {
    a->nodes[ni].b = TT_ERROR;
  }
  a->headers[ni].next_free = a->free_chain;
  a->free_chain = ni;
  a->live_count--;

  for (i = 0; i < a->max_probe+1; i++) {
    unsigned int index = (h + i) % a->table_length;
    tt_node_idx_t candidate = a->table[index];

    printf("hunting i=%d index=%d ni=%d candidate=%d\n", i, index, ni, candidate);
    assert(candidate >= TT_FIRST_VALID_NODE_IDX); /* Internal error if node not in table */

    if (candidate == ni) {
      /* We found it. Now swap in elements. */
      while (1) {
	unsigned int nextindex = (index + 1) % a->table_length;
	tt_node_idx_t next_n = a->table[nextindex];
	tt_hash_t next_h;
	int distance;

	a->table[index] = TT_ERROR;

	if (next_n < TT_FIRST_VALID_NODE_IDX) {
	  break;
	}

	next_h = tt_hash_node(a, next_n);
	distance = nextindex - (next_h % a->table_length);
	if (distance < 0) distance += a->table_length;

	if (distance == 0) {
	  break;
	}

	a->table[index] = next_n;
	index = nextindex;
      }
      break;
    }
  }
}

tt_node_idx_t tt_arena_cons(tt_arena_t *a,
			    uint32_t tag,
			    uint32_t nindex,
			    tt_node_idx_t na,
			    tt_node_idx_t nb)
{
  tt_hash_t h = hash(tag, nindex, na, nb);
  int i;

  for (i = 0; i < a->max_probe+1; i++) {
    unsigned int index = (h + i) % a->table_length;
    tt_node_idx_t candidate = a->table[index];

    printf("cons at %d candidate %d\n", i, candidate);
    /* TODO: perhaps also bail early if we detect that the hash code changes */
    if (candidate < TT_FIRST_VALID_NODE_IDX) {
      printf("cons empty cell\n");
      break;
    }

    printf("tag %d %d\n", a->headers[candidate].inuse.tag, tag);
    printf("index %d %d\n", a->headers[candidate].inuse.index, nindex);
    printf("a %d %d\n", a->nodes[candidate].a, na);
    printf("b %d %d\n", a->nodes[candidate].b, nb);

    if (a->headers[candidate].inuse.tag == tag &&
	a->headers[candidate].inuse.index == nindex &&
	a->nodes[candidate].a == na &&
	a->nodes[candidate].b == nb) {
      printf("cons located correct candidate\n");
      return candidate;
    }
  }

  printf("cons needs to alloc\n");

  if (a->free_chain == TT_ERROR) {
    if (tt_grow(a) != 0) {
      return TT_ERROR;
    }
  }

  {
    tt_node_idx_t node = a->free_chain;
    tt_node_idx_t tostore = node;

    tt_grab(a, na);
    if (tag != TT_TAG_LEAF) tt_grab(a, nb);

    a->free_chain = a->headers[node].next_free;
    tt_drop(a, a->nodes[node].a);
    tt_drop(a, a->nodes[node].b);
    a->live_count++;

    a->headers[node].inuse.refcount = 0;
    a->headers[node].inuse.tag = tag;
    a->headers[node].inuse.index = nindex;
    a->nodes[node].a = na;
    a->nodes[node].b = nb;

    /* Not found */
    i = 0;
    while (1) {
      unsigned int index = (h + i) % a->table_length;
      tt_node_idx_t candidate = a->table[index];

      printf("checking robinhood at h %d i %d index %d candidate %d\n", h, i, index, candidate);

      if (i > a->max_probe) {
	a->max_probe = i;
      }

      if (candidate < TT_FIRST_VALID_NODE_IDX) {
	/* This slot in the table is free. */
	printf("slot free!\n");
	a->table[index] = tostore;
	break;
      }

      printf("slot not free.\n");
      {
	tt_hash_t candidate_h = tt_hash_node(a, candidate);
	int distance = index - (candidate_h % a->table_length);
	if (distance < 0) distance += a->table_length;

	if (distance < i) {
	  a->table[index] = tostore;
	  h = candidate_h;
	  i = distance + 1;
	  tostore = candidate;
	} else {
	  /* keep scanning. */
	  i++;
	}
      }
    }

    return node;
  }
}

tt_node_idx_t tt_grab(tt_arena_t *a, tt_node_idx_t i) {
  if (i >= TT_FIRST_VALID_NODE_IDX && a->headers[i].inuse.refcount < TT_REFCOUNT_LIMIT) {
    a->headers[i].inuse.refcount++;
  }
  return i;
}

void tt_drop(tt_arena_t *a, tt_node_idx_t i) {
  if (i >= TT_FIRST_VALID_NODE_IDX && a->headers[i].inuse.refcount < TT_REFCOUNT_LIMIT) {
    printf("++++++++++++++++++++++++++++++ dropping %d\n", i);
    if (--(a->headers[i].inuse.refcount) == 0) {
      recycle_node(a, i);
    }
  }
}