#include "eval.h"
#include "builtins.h"
#include "env.h"
#include "syntax.h"
#include "util.h"
#include "value.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>



typedef union amb_ptr {
  expr_t *unevaled;
  value_t *evaled;
} amb_ptr;


typedef struct {
  enum {
    FRAME_CALL_ARGS,
    FRAME_CALL_CALLEE,
    FRAME_CALL_ARGS_GROUP,
    FRAME_CALL_ARGS_TUPLE,
    FRAME_ALT,
    FRAME_FGROUP_ALT,
    FRAME_SEQ,
    FRAME_LET,
    FRAME_ENV
  } tag;
  source_t *source;
  pos_t pos;
  union {
    struct {
      expr_t **arguments;
      size_t nargs;
    } call_callee;
    struct {
      expr_t **arguments;
      size_t nargs;
      size_t arg_base;
    } call_args;
    struct {
      expr_t *orelse;
    } alt;
    struct {
      value_t *orelse;
    } fgroup_alt;
    struct {
      expr_t *after;
    } seq;
    struct {
      pattern_t pattern;
      expr_t *in;
    } let_binding;
    struct {
      size_t locals_top;
      size_t locals_frame; 
      value_t *captures; /* should be an env value */
    } env;
  } as;  
} frame_t;
typedef enum {
  MODE_EVAL, MODE_RETURN
} eval_mode_t;

void _frame_debug_dump(frame_t *it) {
  switch (it->tag) {
    case FRAME_ALT:
      printf("ALT(");
      syntax_dump_expr(it->as.alt.orelse);
      printf(")");
      break;
    case FRAME_SEQ:
      printf("SEQ(");
      syntax_dump_expr(it->as.seq.after);
      printf(")");
      break;
    case FRAME_FGROUP_ALT:
      printf("ALTF(");
      value_debug_dump(it->as.fgroup_alt.orelse);
      printf("..?)");
      break;
    
    case FRAME_ENV:
      printf("ENV ..");
      break;
    case FRAME_LET:
      printf("LET");
      break;
    case FRAME_CALL_ARGS_GROUP:
    case FRAME_CALL_ARGS_TUPLE:
      printf("CAT");
      break;
    case FRAME_CALL_ARGS:
      printf("CA");
      break;
    case FRAME_CALL_CALLEE:
      printf("CC");
      break;        
  }
}
static inline size_t stack_bump(frame_t **stack, size_t stack_size, size_t stack_capacity) {
  if (stack_size == stack_capacity) {
    stack_capacity *= 2;
    *stack = realloc(*stack, stack_capacity);
    if (*stack == NULL)
      die("Out of memory: cannot enlarge stack");
  }
  return stack_capacity;
}

static inline size_t arg_stack_bump(value_t ***stack, size_t stack_size, size_t stack_capacity) {
  if (stack_size == stack_capacity) {
    stack_capacity *= 2;
    *stack = realloc(*stack, stack_capacity);
    if (*stack == NULL)
      die("Out of memory: cannot enlarge arg stack");
  }
  return stack_capacity;
}

bool _pattern_match(value_t *value,value_t ***bindings, expr_t *pattern) {
  switch (pattern->tag) {
    case EXPR_CALL:
      if (pattern->as.call.callee->tag == EXPR_CONSTRUCTOR
       && value->tag == VALUE_CONSTRUCTOR 
       && pattern->as.call.nargs == value->as.constructor.num_args
       && pattern->as.call.callee->as.constructor == value->as.constructor.name) {
        for (int i = 0; i < pattern->as.call.nargs; i++) {
          if (!_pattern_match(value->as.constructor.args[i], bindings, pattern->as.call.arguments[i])) {
            return false;
          }
        }
        return true;
      } else return false;
    case EXPR_TUPLE:
      // TODO
      return false;
    case EXPR_CONSTRUCTOR:
      return (value->tag == VALUE_CONSTRUCTOR 
       && pattern->as.call.nargs == 0
       && pattern->as.constructor == value->as.constructor.name);
    case EXPR_PAT_IDENT: {
      printf("%s:",pattern->as.ident);
      (*bindings)[0] = value;
      value_debug_dump(value);
      (*bindings)++;
      printf("\n");
      return true;
    }
    case EXPR_INT_LIT: return (value->tag == VALUE_INT && value->as.integer == pattern->as.int_lit);
    case EXPR_FLOAT_LIT: return (value->tag == VALUE_FLOAT && value->as.floating == pattern->as.float_lit);
    case EXPR_STRING_LIT: return (value->tag == VALUE_STRING && strcmp(value->as.string.string, pattern->as.string_lit) == 0);
    default:
      break;
  }
  return false;
}

value_t *eval(expr_t *expr, gc_t *gc) {

  frame_t *stack = malloc(sizeof(frame_t) * 32); 
  if (stack == NULL)
    die("Out of memory: cannot allocate control stack");
  value_t** locals = malloc(sizeof(value_t *) * 64); 
  if (locals == NULL)
    die("Out of memory: cannot allocate locals stack");
  value_t** arguments = malloc(sizeof(value_t *) * 64); 
  if (arguments == NULL)
    die("Out of memory: cannot allocate arguments stack");
  
  size_t stack_size = 0;
  size_t stack_capacity = 32;
  size_t locals_capacity = 64;
  size_t locals_frame = 0;
  size_t locals_size = 0;
  
  size_t arguments_capacity = 64;
  size_t arguments_size = 0;
  
  eval_mode_t mode = MODE_EVAL;
  amb_ptr focus;
  focus.unevaled = expr;
  value_t *env_val = value_alloc(gc);
  env_val->tag = VALUE_ENV;
  env_val->as.env = NULL;
  while (!(mode == MODE_RETURN && stack_size == 0)) {
    /* debugging code: */
    for (int i = 0; i < stack_size; i++) {
      _frame_debug_dump(stack + i);
      if (i < stack_size - 1)
        printf(", ");
    }
    printf(" | ");
    for (int i = 0; i < locals_size; i++) {
      if (i == locals_size - locals_frame - 1) {
        printf("[");
      }
      value_debug_dump(locals[i]);
      if (i == locals_size - locals_frame - 1) {
        printf("]");
      }
      if (i < locals_size - 1)
        printf(", ");
    }
    printf("{");
    if (env_val->as.env) {env_debug_dump(env_val->as.env);};
    printf("}");
    printf(" | ");
    for (int i = 0; i < arguments_size; i++) {
      value_debug_dump(arguments[i]);
      if (i < stack_size - 1)
        printf(", ");
    }
    if (mode == MODE_EVAL) {
      printf(" > ");
      syntax_dump_expr(focus.unevaled);
      printf("\n");
    } else {
      printf(" < ");
      value_debug_dump(focus.evaled);
      printf("\n");
    }
    
    switch (mode) {
      case MODE_EVAL: switch (focus.unevaled->tag) {
        case EXPR_VAR: {
          value_t *val = (focus.unevaled->as.var >= locals_frame) ? 
            env_lookup(env_val->as.env, focus.unevaled->as.var - locals_frame) :
            locals[locals_size - 1 - focus.unevaled->as.var];
          mode = MODE_RETURN;
          focus.evaled = val;
          break;
        }
        case EXPR_BUILTIN: {
          builtin_func_t *func = focus.unevaled->as.builtin;
          value_t * val = value_alloc(gc);
          val->tag = VALUE_BUILTIN;
          val->as.builtin = func;
          focus.evaled = val;
          mode = MODE_RETURN;
          break;
        }
        case EXPR_INT_LIT: {
          focus.evaled = value_alloc_int(gc,focus.unevaled->as.int_lit);
          mode = MODE_RETURN;
          break;
        }
        case EXPR_FLOAT_LIT: {
          focus.evaled = value_alloc_int(gc,focus.unevaled->as.float_lit);
          mode = MODE_RETURN;
          break;
        }
        case EXPR_STRING_LIT: {
          focus.evaled = value_alloc_string(gc,focus.unevaled->as.string_lit,true);
          mode = MODE_RETURN;
          break;
        }
        case EXPR_CONSTRUCTOR: {
          value_t *val = value_alloc(gc);
          val->tag = VALUE_CONSTRUCTOR;
          val->as.constructor.args = NULL;
          val->as.constructor.num_args = 0;
          val->as.constructor.name = focus.unevaled->as.constructor;
          focus.evaled = val;
          mode = MODE_RETURN;
          break;
        }
        case EXPR_LAMBDA_COMPILED: {
          value_t *val = value_alloc(gc);
          val->tag = VALUE_CLOSURE;
          val->as.closure.body = focus.unevaled->as.lambda.body;
          val->as.closure.env = value_alloc(gc);
          env_t *captured = env_capture(env_val->as.env, locals, locals_frame,locals_size, focus.unevaled->as.lambda.mask);
          val->as.closure.env->tag = VALUE_ENV;
          val->as.closure.env->as.env = captured;
          val->as.closure.next = NULL;
          val->as.closure.arity = focus.unevaled->as.lambda.args;
          focus.evaled = val;
          mode = MODE_RETURN;
          break;
        }
        case EXPR_ALT:
          stack_capacity = stack_bump(&stack, stack_size, stack_capacity);
          stack[stack_size].as.alt.orelse = focus.unevaled->as.alt.orelse;
          stack[stack_size].tag = FRAME_ALT;
          stack_size++;
          focus.unevaled = focus.unevaled->as.alt.try;
          break;
        case EXPR_SEQ:
          stack_capacity = stack_bump(&stack, stack_size, stack_capacity);
          stack[stack_size].as.seq.after = focus.unevaled->as.seq.then;
          stack[stack_size].tag = FRAME_SEQ;
          focus.unevaled = focus.unevaled->as.seq.first;
          stack_size++;  
          break;
        case EXPR_GROUP:
          stack_capacity = stack_bump(&stack, stack_size, stack_capacity);
          stack[stack_size].as.call_args.arguments = focus.unevaled->as.group.clauses + 1;
          stack[stack_size].as.call_args.nargs = focus.unevaled->as.group.nclauses - 1;
          stack[stack_size].as.call_args.arg_base = arguments_size;
          stack[stack_size].tag = FRAME_CALL_ARGS_GROUP;
          stack_size++;
          focus.unevaled = focus.unevaled->as.group.clauses[0];
          break;
        case EXPR_CALL:
          stack_capacity = stack_bump(&stack, stack_size, stack_capacity);
          stack[stack_size].as.call_callee.arguments = focus.unevaled->as.call.arguments;
          stack[stack_size].as.call_callee.nargs = focus.unevaled->as.call.nargs;
          stack[stack_size].tag = FRAME_CALL_CALLEE;
          stack_size++;
          focus.unevaled = focus.unevaled->as.call.callee;
          break;
        case EXPR_LET_BINDING:
          stack_capacity = stack_bump(&stack, stack_size, stack_capacity);
          stack[stack_size].as.let_binding.in = focus.unevaled->as.let_binding.in;
          stack[stack_size].as.let_binding.pattern = focus.unevaled->as.let_binding.pattern;
          stack[stack_size].tag = FRAME_LET;
          stack_size++;
          focus.unevaled = focus.unevaled->as.let_binding.expr;
          break;
        case EXPR_PAT_IDENT:
        case EXPR_LAMBDA:
          die("Impossible: uncompiled expression found while evaluating");
      } 
      break;
      case MODE_RETURN: switch (stack[--stack_size].tag) {
        case FRAME_ENV:
          env_val = stack[stack_size].as.env.captures;
          locals_frame = stack[stack_size].as.env.locals_frame;
          locals_size = stack[stack_size].as.env.locals_top;
          break;
        case FRAME_ALT:
          if (focus.evaled == NULL) {
            focus.unevaled = stack[stack_size].as.alt.orelse;
            mode = MODE_EVAL;
          }
          break;
        case FRAME_FGROUP_ALT:
          if (focus.evaled == NULL) {
            env_val = stack[stack_size].as.fgroup_alt.orelse->as.closure.env;
            focus.unevaled = stack[stack_size].as.fgroup_alt.orelse->as.closure.body;
            mode = MODE_EVAL;
          }
          break;        
        case FRAME_SEQ:
          if (focus.evaled != NULL) {
            focus.unevaled = stack[stack_size].as.seq.after;
            mode = MODE_EVAL;
          }
          break;
        case FRAME_LET: {
          size_t bound_vars = stack[stack_size].as.let_binding.pattern.bound_vars;
          locals_capacity = arg_stack_bump(&locals, locals_size + bound_vars - 1, locals_capacity);
          value_t **write_to = locals + locals_size;
          if (_pattern_match(focus.evaled,&write_to, stack[stack_size].as.let_binding.pattern.expr)) {
            locals_size += bound_vars;
            locals_frame += bound_vars;
            focus.unevaled = stack[stack_size].as.let_binding.in;
            mode = MODE_EVAL;          
            if (stack_size > 0 && stack[stack_size-1].tag != FRAME_ENV || stack_size == 0) {
              stack[stack_size].tag = FRAME_ENV;
              stack[stack_size].as.env.captures = env_val;
              stack[stack_size].as.env.locals_top = locals_size - 1;
              stack[stack_size].as.env.locals_frame = locals_frame;
              stack_size++;
            }
          } else {
            focus.evaled = NULL;
          }
          break;
        } 
        case FRAME_CALL_CALLEE: {          
          expr_t ** args = stack[stack_size].as.call_callee.arguments;
          size_t nargs = stack[stack_size].as.call_callee.nargs;
          stack[stack_size].tag = FRAME_CALL_ARGS;
          stack[stack_size].as.call_args.arg_base = arguments_size;
          stack[stack_size].as.call_args.arguments = args;
          stack[stack_size].as.call_args.nargs = nargs;
          // fall through
        }
        case FRAME_CALL_ARGS_GROUP:
        case FRAME_CALL_ARGS: {
          arguments_capacity = arg_stack_bump(&arguments, arguments_size, arguments_capacity);
          arguments[arguments_size] = focus.evaled;
          arguments_size++;
          if (stack[stack_size].as.call_args.nargs > 0) {
            focus.unevaled = stack[stack_size].as.call_args.arguments[0];
            stack[stack_size].as.call_args.nargs--;
            stack[stack_size].as.call_args.arguments++;
            mode = MODE_EVAL;
            stack_size++;
          } else if (stack[stack_size].tag == FRAME_CALL_ARGS_GROUP) {
            size_t num_args = arguments_size - stack[stack_size].as.call_args.arg_base;
            value_t **args_start = arguments + stack[stack_size].as.call_args.arg_base;            
            printf("num_args %zu arg_base %zu\n",num_args,stack[stack_size].as.call_args.arg_base);
            arguments_size = stack[stack_size].as.call_args.arg_base;
            value_t *ret = args_start[num_args-1];
            if (ret->tag != VALUE_CLOSURE) {
              focus.evaled = NULL;
              break;
            }
            size_t i = num_args;
            while (i >= 2) {
              if (args_start[i-2]->tag != VALUE_CLOSURE || args_start[i-2]->as.closure.arity != ret->as.closure.arity) {
                ret = NULL;
                break;
              }
              value_t *nret = value_alloc(gc);
              nret->tag = VALUE_CLOSURE;
              nret->as.closure.arity = ret->as.closure.arity;
              nret->as.closure.env = args_start[i-2]->as.closure.env;
              nret->as.closure.body = args_start[i-2]->as.closure.body;
              nret->as.closure.next = ret;
              ret = nret;   
              i--;
            }
            focus.evaled = ret;
          } else if (stack[stack_size].tag == FRAME_CALL_ARGS) { 
            value_t *callee = arguments[stack[stack_size].as.call_args.arg_base];
            size_t num_args = arguments_size - stack[stack_size].as.call_args.arg_base - 1;
            value_t **args_start = arguments + stack[stack_size].as.call_args.arg_base + 1;
            arguments_size = stack[stack_size].as.call_args.arg_base;
            printf("num_args %zu arg_base %zu\n",num_args,stack[stack_size].as.call_args.arg_base);
            printf("callee %d\n",callee->as.closure.arity);
            if (callee == NULL) {
              focus.evaled = NULL;
            } else switch (callee->tag) {
              case VALUE_BUILTIN:
                focus.evaled = callee->as.builtin(gc,args_start, num_args);
                break;
              case VALUE_CLOSURE:                
                if (callee->as.closure.arity == num_args) {
                  printf("ARITY CHECK\n");
                  if (stack_size > 0 && stack[stack_size-1].tag != FRAME_ENV || stack_size == 0) {
                    stack[stack_size].as.env.captures = env_val;
                    stack[stack_size].as.env.locals_frame = locals_frame;
                    stack[stack_size].as.env.locals_top = locals_size;
                    stack[stack_size].tag = FRAME_ENV;
                    stack_size++;
                  }                  
                  locals_frame = num_args;
                  locals_capacity = arg_stack_bump(&locals, locals_size + num_args, locals_capacity);
                  memcpy(locals + locals_size, args_start, sizeof(value_t*) * num_args);
                  locals_size += num_args;
                  
                  env_val = callee->as.closure.env;
                  focus.unevaled = callee->as.closure.body;
                  if (callee->as.closure.next != NULL) {
                    stack_bump(&stack, stack_size, stack_capacity);
                    stack[stack_size].as.fgroup_alt.orelse = callee->as.closure.next;
                    stack[stack_size].tag = FRAME_FGROUP_ALT;
                    stack_size++;
                  }
                  mode = MODE_EVAL;
                } else {
                  focus.evaled = NULL;
                }
                break;
              case VALUE_CONSTRUCTOR:
                if (callee->as.constructor.args == NULL) {
                  value_t *val = value_alloc(gc);
                  val->tag = VALUE_CONSTRUCTOR;
                  val->as.constructor.args = malloc(sizeof(value_t*) * num_args);
                  if (val->as.constructor.args == NULL) {
                    die("Out of memory, can't allocate args buffer");
                  }
                  memcpy(val->as.constructor.args,args_start, sizeof(value_t*) * num_args);
                  val->as.constructor.num_args = num_args;
                  val->as.constructor.name = callee->as.constructor.name;
                  focus.evaled = val;
                  break;
                }
                //fall through  
              case VALUE_INT:
              case VALUE_FLOAT:
              case VALUE_STRING:              
                focus.evaled = NULL;
              case VALUE_ENV: // should be impossible
                focus.evaled = NULL;
            }
          }
          break;
        }
      }
      break;
    }
  }
  free(arguments);
  free(locals);
  free(stack);
  return focus.evaled;
}