#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "expression.h"


struct li_production_t * li_production_alloc(size_t seq_len) {
  struct li_production_t *prod = malloc(sizeof(struct li_production_t));
  if (prod == NULL) {
    fprintf(stderr, "failed to allocate production rule\n");
    return NULL;
  }
  prod->seq = malloc(seq_len * sizeof(int));
  if (prod->seq == NULL) {
    fprintf(stderr, "failed to allocate production sequence buffer\n");
    free(prod);
    return NULL;
  }
  prod->seq_len = seq_len;
  prod->next = NULL;
  return prod;
}


void li_production_free(struct li_production_t *prod) {
  if (prod != NULL) {
    li_production_free(prod->next);
    free(prod->seq);
    free(prod);
  }
}


struct li_production_rule_t * li_prod_rule_alloc() {
  struct li_production_rule_t *rule = malloc(sizeof(struct li_production_rule_t));
  if (rule == NULL) {
    fprintf(stderr, "failed to allocate production rule\n");
    return NULL;
}
  rule->num_productions = 0;
  rule->head = NULL;
  return rule;
}


void li_prod_rule_free(struct li_production_rule_t *rule) {
  if (rule != NULL) {
    li_production_free(rule->head);
    free(rule);
  }
}


int li_prod_rule_add(
  struct li_production_rule_t *rule, size_t seq_len, int *seq
) {
  struct li_production_t *prod = li_production_alloc(seq_len);
  if (prod == NULL) {
    return 1;
  }
  memcpy(prod->seq, seq, seq_len*sizeof(int));
  prod->next = rule->head;
  rule->head = prod;
  rule->num_productions += 1;
  return 0;
}


li_grammar_t * li_grammar_alloc(int num_nonterminal) {
  li_grammar_t *grammar = malloc(sizeof(li_grammar_t));
  if (grammar == NULL) {
    fprintf(stderr, "failed to allocate grammar\n");
    return NULL;
  }
  grammar->rules = malloc(num_nonterminal * sizeof(struct li_production_rule_t));
  if (grammar->rules == NULL) {
    fprintf(stderr, "failed to allocate grammar rules buffer\n");
    free(grammar);
    return NULL;
  }
  for (int i=0; i<num_nonterminal; i++) {
    grammar->rules[i].num_productions = 0;
    grammar->rules[i].head = NULL;
  }
  grammar->num_nonterminal = num_nonterminal;
  return grammar;
}


void li_grammar_free(li_grammar_t *grammar) {
  if (grammar != NULL) {
    for (int i=0; i<grammar->num_nonterminal; i++) {
      li_production_free(grammar->rules[i].head);
    }
    free(grammar->rules);
    free(grammar);
  }
}


int li_grammar_rule_add(
  li_grammar_t *grammar, int nonterminal, size_t seq_len, int *seq
) {
  if (nonterminal >= grammar->num_nonterminal) {
    return 2;
  }
  return li_prod_rule_add(grammar->rules + nonterminal, seq_len, seq);
}



int stack_push(int *stack, size_t stack_sz, size_t *top, int value) {
  *top += 1;
  if (*top >= stack_sz) {
    fprintf(stderr, "grammar stack overflow!\n");
    return 1;
  }
  stack[*top] = value;
  return 0;
}


int stack_pop(int *stack, size_t *top, int *value) {
  if (*top == ((size_t)-1)) {
    fprintf(stderr, "grammar stack underflow!\n");
    return 1;
  }

  *value = stack[*top];
  *top -=1;
  return 0;
}


int expand_production(
  li_grammar_t *grammar, 
  int *dst, size_t sz, size_t *len,
  int *stack, size_t stack_sz, size_t *top,
  int symbol,
  const uint8_t *randomness, size_t rand_sz
) {
  struct li_production_rule_t rule = grammar->rules[symbol];
  struct li_production_t *prod;
  if (rule.num_productions == 0) {
    // empty rule, nothing to do
    return 0;
  } else if (rule.num_productions == 1) {
    // pick only production
    prod = rule.head;
  } else {
    // use randomness to select production
    int selection = randomness[0] % rule.num_productions;
    randomness += 1;
    rand_sz -= 1;
    prod = rule.head;
    for (int i=0; i<selection; i++) {
      prod = prod->next;
    }
  }

  // push production to stack
  for (int i=prod->seq_len-1; i>=0; i--) {
    if (stack_push(stack, stack_sz, top, prod->seq[i]) != 0) {
      return 1;
    }
  }

  // pop production from stack, recursing on nonterminals
  while (*top != ((size_t)-1)) {
    int symbol;
    if (stack_pop(stack, top, &symbol) != 0) {
      return 1;
    }
    if (symbol < grammar->num_nonterminal) {
      int err = expand_production(grammar, dst, sz, len, stack, stack_sz, top, symbol, randomness, rand_sz);
      if (err != 0) {
        return err;
      }
    } else {
      if (*len > sz) {
        fprintf(stderr, "exceeded maximum output buffer length\n");
        return 1;
      }
      dst[*len] = symbol;
      *len += 1;
    }
  }

  return 0;
}



int li_grammar_rule_expand(
  li_grammar_t *grammar, 
  int *dst, size_t *len, 
  int *stack, size_t stack_sz, 
  int symbol, 
  const uint8_t *randomness, size_t rand_sz
) {
  size_t sz = *len;
  *len = 0;
  size_t top = -1;
  return expand_production(grammar, dst, sz, len, stack, stack_sz, &top, symbol, randomness, rand_sz);
}