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

#include "treetrie.h"
#include "critbit.h"
#include "route.h"

static inline tt_node_ptr_t rseq(tt_arena_t *a, tt_atom_t s, tt_node_ptr_t r) {
  if (TT_EMPTY_P(r)) {
    return TT_EMPTY;
  }
  assert(s != TT_WILD);
  return tt_cons_branch(a, TT_EMPTY, RET_IF_NO_PTR(tt_dict_singleton(a, s, r)));
}

static inline tt_node_ptr_t rwild(tt_arena_t *a, tt_node_ptr_t r) {
  if (TT_EMPTY_P(r)) {
    return TT_EMPTY;
  }
  return tt_cons_branch(a, r, TT_EMPTY_DICT);
}

static inline tt_node_ptr_t rbranch(tt_arena_t *a,
				    tt_node_ptr_t wild, /* trie */
				    tt_node_ptr_t others) /* dict */
{
  if (TT_EMPTY_P(wild) && TT_EMPTY_DICT_P(others)) {
    return TT_EMPTY;
  }
  return tt_cons_branch(a, wild, others);
}

static inline tt_node_ptr_t rwildseq(tt_arena_t *a, tt_node_ptr_t r) {
  if (TT_EMPTY_P(r)) {
    return TT_EMPTY;
  }
  return tt_cons_tail(a, r);
}

static inline tt_node_ptr_t runwildseq(tt_arena_t *a, tt_node_ptr_t r) {
  if (tt_ptr_tag(r) == TT_TAG_TAIL) {
    return TT_TAIL_TRIE(a,r);
  } else {
    return TT_EMPTY;
  }
}

static inline tt_node_ptr_t rlookup_dict(tt_arena_t *a,
					 tt_node_ptr_t wild,
					 tt_node_ptr_t others,
					 tt_atom_t key)
{
  tt_node_ptr_t result;
  result = tt_dict_get(a, others, key);
  if (!TT_NO_PTR_P(result)) {
    return result;
  }
  result = wild;
  switch (key) {
    case TT_BOS: return rwildseq(a, result);
    case TT_EOS: return runwildseq(a, result);
    default:     return result;
  }
}

static inline tt_node_ptr_t rlookup(tt_arena_t *a,
				    tt_node_ptr_t r,
				    tt_atom_t key)
{
  assert(tt_ptr_tag(r) == TT_TAG_BRANCH);
  return rlookup_dict(a, TT_BRANCH_WILDCARD(a,r), TT_BRANCH_OTHERS(a,r), key);
}

static inline tt_node_ptr_t rupdate_dict(tt_arena_t *a,
					 tt_node_ptr_t old_wild, /* trie */
					 tt_node_ptr_t old_others, /* dict */
					 tt_atom_t key,
					 tt_node_ptr_t k) /* trie */
{
  assert(key != TT_WILD);
  assert(TT_EMPTY_DICT_P(old_others) || tt_ptr_tag(old_others) == TT_TAG_DICT);

#define OTHERS_SANS_KEY() (tt_dict_remove(a, old_others, key))
#define OTHERS_WITH_KEY() (tt_dict_set(a, old_others, key, k))

  switch (key) {
    case TT_BOS:
      if (tt_ptr_tag(k) == TT_TAG_TAIL) {
	return (TT_TAIL_TRIE(a,k) == old_wild) ? OTHERS_SANS_KEY() : OTHERS_WITH_KEY();
      } else {
	return TT_EMPTY_P(k) ? OTHERS_SANS_KEY() : OTHERS_WITH_KEY();
      }
    case TT_EOS:
      if (tt_ptr_tag(old_wild) == TT_TAG_TAIL) {
	return (TT_TAIL_TRIE(a,old_wild) == k) ? OTHERS_SANS_KEY() : OTHERS_WITH_KEY();
      } else {
	return TT_EMPTY_P(k) ? OTHERS_SANS_KEY() : OTHERS_WITH_KEY();
      }
    default:
      return (k == old_wild) ? OTHERS_SANS_KEY() : OTHERS_WITH_KEY();
  }

#undef OTHERS_SANS_KEY
#undef OTHERS_WITH_KEY
}

static inline tt_node_ptr_t rupdate(tt_arena_t *a,
				    tt_node_ptr_t r0, /* branch */
				    tt_atom_t key,
				    tt_node_ptr_t k) /* trie */
{
  assert(key != TT_WILD);
  if (TT_EMPTY_P(r0)) {
    return rseq(a, key, k);
  } else {
    tt_node_ptr_t new_others;
    assert(tt_ptr_tag(r0) == TT_TAG_BRANCH);
    new_others = rupdate_dict(a, TT_BRANCH_WILDCARD(a,r0), TT_BRANCH_OTHERS(a,r0), key, k);
    return rbranch(a, TT_BRANCH_WILDCARD(a,r0), RET_IF_NO_PTR(new_others));
  }
}

static inline tt_node_ptr_t expand(tt_arena_t *a, tt_node_ptr_t tailnode) {
  tt_node_ptr_t others = RET_IF_NO_PTR(tt_dict_singleton(a, TT_EOS, TT_TAIL_TRIE(a,tailnode)));
  return rbranch(a, tailnode, others);
}

/* N.B. Returns a tt_grab'd result. */
tt_node_ptr_t tt_trie_combine(tt_arena_t *a,
			      tt_node_ptr_t r1,
			      tt_node_ptr_t r2,
			      int left_empty_keep,
			      int right_empty_keep,
			      int left_base_keep,
			      int right_base_keep,
			      void *f_context,
			      /* Should return a tt_grab'd result. */
			      tt_node_ptr_t (*f)(void *f_context,
						 tt_node_ptr_t r1,
						 tt_node_ptr_t r2))
{
  /* N.B. Returns a tt_grab'd result. */
  tt_node_ptr_t g(tt_node_ptr_t r1, tt_node_ptr_t r2) {
    tt_tag_t t1, t2;

    if (TT_EMPTY_P(r1)) {
      return left_empty_keep ? tt_grab(a, r2) : TT_EMPTY;
    }
    if (TT_EMPTY_P(r2)) {
      return right_empty_keep ? tt_grab(a, r1) : TT_EMPTY;
    }

    t1 = tt_ptr_tag(r1);
    t2 = tt_ptr_tag(r2);

    if (t1 == TT_TAG_BRANCH && t2 == TT_TAG_BRANCH) {
      tt_node_ptr_t w1 = TT_BRANCH_WILDCARD(a,r1);
      tt_node_ptr_t w2 = TT_BRANCH_WILDCARD(a,r2);
      tt_node_ptr_t dict1 = TT_BRANCH_OTHERS(a,r1);
      tt_node_ptr_t dict2 = TT_BRANCH_OTHERS(a,r2);
      tt_node_ptr_t result_wild = TT_NO_PTR; /* grab'd */
      tt_node_ptr_t result_others = TT_EMPTY_DICT; /* grab'd */
      tt_node_ptr_t result = TT_NO_PTR; /* grab'd */
      int ok = 1;

      int examine_key(void *examine_key_context, tt_atom_t key, tt_node_ptr_t ignored_trie) {
	tt_node_ptr_t trie1, trie2; /* Looked up in r1 and r2 */
	tt_node_ptr_t new_trie; /* Combination of trie1 and trie2 */
	tt_node_ptr_t new_result_others; /* Updated dictionary with key---new_trie association */
	trie1 = tt_grab(a, rlookup_dict(a, w1, dict1, key));
	if (TT_NO_PTR_P(trie1)) return 0;
	trie2 = tt_grab(a, rlookup_dict(a, w2, dict2, key));
	if (TT_NO_PTR_P(trie2)) {
	  tt_drop(a, trie1);
	  return 0;
	}
	new_trie = g(trie1, trie2); /* already grabbed */
	tt_drop(a, trie1);
	tt_drop(a, trie2);
	if (TT_NO_PTR_P(new_trie)) return 0;
	new_result_others = tt_grab(a, rupdate_dict(a, result_wild, result_others, key, new_trie));
	tt_drop(a, new_trie);
	if (TT_NO_PTR_P(new_result_others)) return 0;
	tt_drop(a, result_others);
	result_others = new_result_others;
	return 1;
      }

      if (!TT_EMPTY_P(w1) && !TT_EMPTY_P(w2)) {
	/* Two wildcards - worst case. Must loop over both dictionaries. */
	result_wild = g(w1, w2);
	ok = ok && tt_dict_foreach(a, dict1, NULL, examine_key);
	ok = ok && tt_dict_foreach(a, dict2, NULL, examine_key);
      } else if ((!TT_EMPTY_P(w1)) ||
		 (TT_EMPTY_P(w2) &&
		  (tt_dict_size(a, dict1) >= tt_dict_size(a, dict2)))) {
	/* Either a wildcard on the left, or no wildcard at all but the left is larger */
	result_wild = tt_grab(a, w1);
	result_others = left_base_keep ? tt_grab(a, dict1) : TT_EMPTY_DICT;
	ok = ok && tt_dict_foreach(a, dict2, NULL, examine_key);
      } else {
	/* Either a wildcard on the right, or no wildcard at all but the right is larger */
	result_wild = tt_grab(a, w2);
	result_others = right_base_keep ? tt_grab(a, dict2) : TT_EMPTY_DICT;
	ok = ok && tt_dict_foreach(a, dict1, NULL, examine_key);
      }

      if (ok) {
	result = tt_grab(a, rbranch(a, result_wild, result_others));
      }

      tt_drop(a, result_wild);
      tt_drop(a, result_others);
      return result;
    }

    if (t1 == TT_TAG_TAIL) {
      if (t2 == TT_TAG_TAIL) {
	tt_node_ptr_t combined = RET_IF_NO_PTR(g(TT_TAIL_TRIE(a,r1), TT_TAIL_TRIE(a,r2)));
	tt_node_ptr_t result = tt_grab(a, rwildseq(a, combined));
	tt_drop(a, combined);
	return result;
      } else {
	tt_node_ptr_t r1_expanded = tt_grab(a, RET_IF_NO_PTR(expand(a, r1)));
	tt_node_ptr_t result = g(r1_expanded, r2);
	tt_drop(a, r1_expanded);
	return result;
      }
    } else if (t2 == TT_TAG_TAIL) {
      tt_node_ptr_t r2_expanded = tt_grab(a, RET_IF_NO_PTR(expand(a, r2)));
      tt_node_ptr_t result = g(r1, r2_expanded);
      tt_drop(a, r2_expanded);
      return result;
    }

    if (t1 == TT_TAG_OK || t2 == TT_TAG_OK) {
      return f(f_context, r1, r2);
    }

    /* There is no legitimate combination of tags that should let us get here. */
    assert(0);
  }

  /* No need for tt_grab here - g has already done that for us */
  return g(r1, r2);
}

tt_node_ptr_t tt_trie_union(tt_arena_t *a, tt_node_ptr_t r1, tt_node_ptr_t r2) {
  tt_node_ptr_t f_union(void *f_context, tt_node_ptr_t r1, tt_node_ptr_t r2) {
    tt_tag_t t1 = tt_ptr_tag(r1);
    tt_tag_t t2 = tt_ptr_tag(r2);
    if (t1 != TT_TAG_OK) {
      /* t2 must be ok. */
      return tt_grab(a, r2);
    } else if (t2 != TT_TAG_OK) {
      return tt_grab(a, r1);
    } else {
      tt_node_ptr_t s = RET_IF_NO_PTR(tt_dictset_union(a, TT_OK_DICT(a,r1), TT_OK_DICT(a,r2)));
      tt_node_ptr_t result = tt_cons_ok(a, s);
      tt_drop(a, s);
      return tt_grab(a, result);
    }
  }
  return tt_trie_combine(a, r1, r2, 1, 1, 1, 1, NULL, f_union);
}

tt_node_ptr_t tt_begin_path(tt_arena_t *a, tt_node_ptr_t ok_dict) {
  return tt_cons_ok(a, ok_dict);
}

tt_node_ptr_t tt_prepend_path(tt_arena_t *a, tt_atom_t tok, tt_node_ptr_t tail) {
  if (tok == TT_WILD) {
    return rwild(a, tail);
  } else {
    return rseq(a, tok, tail);
  }
}

void print_indent(int spaces) {
  while (spaces--) {
    putchar(' ');
  }
}

void tt_dump_routingtable(tt_arena_t *a, tt_node_ptr_t r, int initial_indent) {
  void walk(int indent, tt_node_ptr_t r) {
    switch (tt_ptr_tag(r)) {
      case TT_TAG_TAIL:
	printf(" ...>");
	walk(indent + 5, TT_TAIL_TRIE(a, r));
	break;
      case TT_TAG_OK: {
	int need_space = 0;
	int pkey(void *context, tt_atom_t key, tt_node_ptr_t ignored_trie) {
	  if (need_space) {
	    putchar(' ');
	  } else {
	    need_space = 1;
	  }
	  printf("%d", key);
	  return 1;
	}
	fputs(" {", stdout);
	tt_dict_foreach(a, TT_OK_DICT(a, r), NULL, pkey);
	putchar('}');
	break;
      }
      case TT_TAG_BRANCH: {
	int need_sep = 0;
	int pedge(void *context, tt_atom_t key, tt_node_ptr_t node) {
	  char keystr[256]; /* not very tight */
	  if (need_sep) {
	    putchar('\n');
	    print_indent(indent);
	  } else {
	    need_sep = 1;
	  }
	  keystr[sizeof(keystr) - 1] = '\0';
	  switch (key) {
	    case TT_WILD: snprintf(keystr, sizeof(keystr) - 1, " ★"); break;
	    case TT_BOS: snprintf(keystr, sizeof(keystr) - 1, " <"); break;
	    case TT_EOS: snprintf(keystr, sizeof(keystr) - 1, " >"); break;
	    case TT_BOC: snprintf(keystr, sizeof(keystr) - 1, " {"); break;
	    case TT_EOC: snprintf(keystr, sizeof(keystr) - 1, " }"); break;
	    default:
	      snprintf(keystr, sizeof(keystr) - 1, " %d", key);
	  }
	  fputs(keystr, stdout);
	  walk(indent + strlen(keystr), node);
	  return 1;
	}
	if (!TT_EMPTY_P(TT_BRANCH_WILDCARD(a, r))) {
	  pedge(NULL, TT_WILD, TT_BRANCH_WILDCARD(a, r));
	}
	tt_dict_foreach(a, TT_BRANCH_OTHERS(a, r), NULL, pedge);
	if (!need_sep) {
	  printf(" ::: no edges!");
	}
	break;
      }
      case TT_TAG_SPECIAL:
	if (TT_EMPTY_P(r)) {
	  printf(" ::: nothing");
	  break;
	}
	/* fall through */
      default:
	printf("?!?!?! %x", (unsigned int) r);
	break;
    }
  }
  walk(initial_indent, r);
  putchar('\n');
}