#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <string.h>
#include "nodes.h"
#include "C.tab.h"
#include "types.h"
#include "list.h"

TOKEN* lookup_token(char *s);
int asprintf(char **strp, const char *fmt, ...);

int debug = 0;

/*
 * Existing problems/todo:
 * return from control block, need a "return address" to jump to -- call levels
 * --multiple assignments `int x,y,z = 1`--
 * new block scope for if statements?
 * arguments for functions
 * difference between tilde and semicolon
 * while loops
 * control flow `continue` and `break` linked to return address sort of? call levels
 */

char *named(int t) {
    static char b[100];
    if (isgraph(t) || t==' ') {
      sprintf(b, "%c", t);
      return b;
    }
    switch (t) {
      default: return "???";
    case IDENTIFIER:
      return "id";
    case CONSTANT:
      return "constant";
    case STRING_LITERAL:
      return "string";
    case LE_OP:
      return "<=";
    case GE_OP:
      return ">=";
    case EQ_OP:
      return "==";
    case NE_OP:
      return "!=";
    case EXTERN:
      return "extern";
    case AUTO:
      return "auto";
    case INT:
      return "int";
    case VOID:
      return "void";
    case APPLY:
      return "apply";
    case LEAF:
      return "leaf";
    case IF:
      return "if";
    case ELSE:
      return "else";
    case WHILE:
      return "while";
    case CONTINUE:
      return "continue";
    case BREAK:
      return "break";
    case RETURN:
      return "return";
    }
}



void print_leaf(NODE *tree, int level) {
    TOKEN *t = (TOKEN *)tree;
    int i;
    for (i=0; i<level; i++) putchar(' ');
    if (t->type == CONSTANT) printf("%d\n", t->value);
    else if (t->type == STRING_LITERAL) printf("\"%s\"\n", t->lexeme);
    else if (t) puts(t->lexeme);
}

void print_tree0(NODE *tree, int level) {
    int i;
    if (tree==NULL) return;
    if (tree->type==LEAF) {
      print_leaf(tree->left, level);
    }
    else {
      for(i=0; i<level; i++) putchar(' ');
      printf("%s\n", named(tree->type));

/*       if (tree->type=='~') { */
/*         for(i=0; i<level+2; i++) putchar(' '); */
/*         printf("%p\n", tree->left); */
/*       } */
/*       else */
        print_tree0(tree->left, level+2);
      print_tree0(tree->right, level+2);
    }
}

// forward declare because it is used in add_var_to_env
BIND* recursive_interpret(NODE*, ENV*);

void print_tree(NODE *tree) {
    print_tree0(tree, 0);
}

int count_args(NODE *tree) { 
  NODE *tmp = tree;
  int a = 1;
  while (tmp != NULL) {
    if (tmp->type != ',') {
      break;
    }
    tmp = tmp->left;
    a++;
  }
  return a;
}

void populate_arg_list(NODE *tree, TOKEN** arr, int arr_len) {
  if (tree->type == '~') {
    for (int i = 0; i < arr_len; i++) {
      if (arr[i] == NULL) {
        TOKEN* name = (TOKEN*)tree->right->left;
        arr[i] = name;
        break;
      }
    }
  } else {
    populate_arg_list(tree->left, arr, arr_len);
    populate_arg_list(tree->right, arr, arr_len);
  }
}

TOKEN** get_argument_list(NODE *tree) {
  int num_args = count_args(tree->left->right->right);
  TOKEN** arr = (TOKEN**) malloc(sizeof(TOKEN*) * num_args);
  populate_arg_list(tree->left->right->right, arr, num_args);
  return arr;
}

void add_function_to_env(NODE *tree, ENV *env_ptr) {
  BIND* existing;
  NODE* func_name = tree->left->right->left->left;
  TOKEN* name_token = (TOKEN*) func_name;
  if ((existing = find_name_in_env(name_token, env_ptr)) == NULL) {
    if (env_ptr->bindings == NULL) {
      env_ptr->bindings = create_list(create_binding(name_token, tree, env_ptr), NULL);
    } else {
      append_list(env_ptr->bindings, create_binding(name_token, tree, env_ptr));
    }
  } else {
    printf("Error: redefinition of function with name %s\n", name_token->lexeme);
  }
}

void populate_val_list(NODE* tree, NODE** arr, int num_args, ENV* env_ptr) {
  if (tree->type == ',') {
    populate_val_list(tree->left, arr, num_args, env_ptr);
    populate_val_list(tree->right, arr, num_args, env_ptr);
  } else {
    for (int i = 0; i < num_args; i++) {
      if (arr[i] == NULL) {
        arr[i] = tree;
        break;
      }
    }
  }
}

NODE** get_val_list(NODE *tree, ENV* env_ptr) {
  int num_args = count_args(tree);
  NODE** arr = (NODE**) malloc(sizeof(NODE*) * num_args);
  populate_val_list(tree, arr, num_args, env_ptr);
  return arr;
}

void bind_arg(NODE* val, TOKEN* arg_name, ENV *env_ptr, ENV* src_env) {
  BIND* result = recursive_interpret(val, src_env);
  if(debug)printf("Binding argument: %s\n", arg_name->lexeme);
  if (env_ptr->bindings == NULL) {
    env_ptr->bindings = create_list(create_binding(arg_name, result->tree, result->env), NULL);
  } else {
    BIND* existing = find_name_in_list(arg_name, env_ptr->bindings);
    if (existing) {
      existing->tree = result->tree;
    } else {
      append_list(env_ptr->bindings, create_binding(arg_name, result->tree, result->env));
    }
  }
}

void add_var_to_env(NODE *tree, ENV *env_ptr) {
  BIND* existing;
  NODE* var_name = tree->left->left;
  TOKEN* name_token = (TOKEN*) var_name;
  if (debug) printf("assigning %s:\n", name_token->lexeme);
  BIND* node = recursive_interpret(tree->right, env_ptr);
  if ((existing = find_name_in_env(name_token, env_ptr)) == NULL) {
    if (env_ptr->bindings == NULL) {
      env_ptr->bindings = create_list(create_binding(name_token, node->tree, node->env), NULL);
    } else {
      append_list(env_ptr->bindings, create_binding(name_token, node->tree, node->env));
    }
  } else {
    existing->tree = node->tree;
  } 
}

TOKEN* get_int_token(int value) {
  TOKEN* res;
  char buffer[20];
  snprintf(buffer, 20, "%d", value);
  res = lookup_token(buffer);
  if (res->type != CONSTANT) res->type = CONSTANT;
  if (res->value != value) res->value = value;
  return res;
}

BIND* anon_binding(NODE *tree) {
  return create_binding(NULL, tree, NULL);
}

BIND* ret_val;
int returning;

BIND* recursive_interpret(NODE *tree, ENV *env_ptr) {
  if (returning == 1) return NULL;
  if (tree==NULL) return NULL;
  if (tree->type==LEAF) {
    if (tree->left->type == CONSTANT) {
      return anon_binding(tree->left);
    } else if (tree->left->type == STRING_LITERAL) {
      printf("Not implemented\n");
      exit(1);
    } else if (tree->left->type == INT || tree->left->type == FUNCTION) {
      // do nothing we dont care about types for now
      return NULL;
    } else {
      // an identifier
      TOKEN* tok = (TOKEN *) tree->left;
      BIND* var_bind = find_name_in_env(tok, env_ptr);
      if (debug) {
        printf("looking up %s\n", tok->lexeme);
      }
      if (var_bind == NULL) {
        printf("Could not find variable %s\n", tok->lexeme);
        exit(1);
      }
      if(debug)printf("I got: ");
      if (var_bind->tree->type == CONSTANT) {
        if(debug){
          print_leaf(var_bind->tree, 0);
        }
      } else {
        if(debug)print_tree(var_bind->tree);
      }
      return var_bind;
      /* if (var_bind->tree->type == CONSTANT) { */
      /*   return var_bind->tree; */
      /* } else { */
      /*   print_tree(var_bind->tree); */
      /*   printf("I GOT A %d from %s\n", var_bind->tree->type, var_bind->name->lexeme); */
      /*   printf("Maybe got a function?\n"); */
      /*   return NULL; */
      /* } */
    }
  }
  if (tree->type=='D') {
    // this is a function definition
    add_function_to_env(tree, env_ptr);
    return NULL;
  }
  if (tree->type=='=') {
    // this is a variable definition
    add_var_to_env(tree, env_ptr);
    return NULL;
  }
  if (tree->type==APPLY) {
    TOKEN* func_name = ((TOKEN*) tree->left->left);
    BIND* func = find_name_in_env(func_name, env_ptr);
    if (func == NULL) {
      printf("Could not find binding for function with name %s\n", func_name->lexeme);
      exit(1);
    }
    ENV* new_func_env = create_new_function_env(func->env);
    if (debug) printf("Calling function %s\n", func_name->lexeme);
    if (tree->right != NULL) {
      TOKEN **arg_list;
      NODE **val_list;
      int exp_args, num_args;
      exp_args = count_args(func->tree->left->right->right);
      num_args = count_args(tree->right);
      arg_list = get_argument_list(func->tree);
      val_list = get_val_list(tree->right, env_ptr);
      if (exp_args == num_args) {
        for (int i = 0; i < num_args; i++) {
          bind_arg(val_list[i], arg_list[i], new_func_env, env_ptr);
        }
      } else {
        printf("Incorrect arguments passed to function %s, expected %d got %d\n", func_name->lexeme, exp_args, num_args);
        exit(1);
      }
    }
    recursive_interpret(func->tree->right, new_func_env);
    if (returning == 1) {
      returning = 0;
      return ret_val;
    } else {
      printf("HELP called a function with no return!!\n");
      exit(1);
    }
  }
  if (tree->type == IF) {
    if (((TOKEN*)recursive_interpret(tree->left, env_ptr)->tree)->value == 1) {
      ENV* scope_env = create_new_function_env(env_ptr);
      if (tree->right->type == ELSE) {
        return recursive_interpret(tree->right->left, scope_env);
      } else {
        return recursive_interpret(tree->right, scope_env);
      }
    } else {
      if (tree->right->type == ELSE) {
        ENV* scope_env = create_new_function_env(env_ptr);
        return recursive_interpret(tree->right->right, scope_env);
      }
      return 0;
    }
  }
  if (tree->type == LE_OP) {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value <= rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == GE_OP) {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value >= rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == EQ_OP) {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    if (debug) printf("Comparing %d == %d\n", lhs->value, rhs->value);
    int result = lhs->value == rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == NE_OP) {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value != rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '>') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value > rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '<') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value < rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '+') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    if (debug) printf("Adding %d and %d\n", lhs->value, rhs->value);
    int result = lhs->value + rhs->value;
    return anon_binding((NODE*) get_int_token(result));
 }
  if (tree->type == '-') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value - rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '*') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value * rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '/') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value / rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '%') {
    TOKEN* lhs = (TOKEN*) recursive_interpret(tree->left, env_ptr)->tree;
    TOKEN* rhs = (TOKEN*) recursive_interpret(tree->right, env_ptr)->tree;
    int result = lhs->value % rhs->value;
    return anon_binding((NODE*) get_int_token(result));
  }
  if (tree->type == '~') {
    if (tree->left->type == INT || tree->left->type == FUNCTION) {
      return recursive_interpret(tree->right, env_ptr);
    } else {
      recursive_interpret(tree->left, env_ptr);
      return recursive_interpret(tree->right, env_ptr);
    }
  }
  if (tree->type == RETURN) {
    ret_val = recursive_interpret(tree->left, env_ptr);
    returning = 1;
  }
  if (tree->type == WHILE) {
    ENV* scope_env = create_new_function_env(env_ptr);
    while(((TOKEN*)recursive_interpret(tree->left, env_ptr)->tree)->value == 1) {
      recursive_interpret(tree->right, scope_env);
    }
    return 0;
  }
  recursive_interpret(tree->left, env_ptr);
  return recursive_interpret(tree->right, env_ptr);
}

ENV* cons_global_env(NODE *tree) {
  ENV* global = create_new_function_env(NULL);
  recursive_interpret(tree, global);
  return global;
}

void interpret_tree(NODE *tree) {
  ENV* global_env = cons_global_env(tree);
  BIND* ref_main = find_name_in_env(lookup_token("main"), global_env);
  if (ref_main == NULL) {
    printf("Could not find main, cannot run!\n");
    exit(1);
  }
  // print ref_main to make sure we really got it
  printf("Located %s, ready to run!\n", ref_main->name->lexeme);
  recursive_interpret(ref_main->tree->right, ref_main->env);
  /* printf("%d\n", ((TOKEN*)recursive_interpret(ref_main->tree->right, ref_main->env))->value); */
  printf("%d\n", ((TOKEN*)ret_val->tree)->value);
}

extern int yydebug;
extern NODE* yyparse(void);
extern NODE* ans;
extern void init_symbtable(void);

int main(int argc, char** argv)
{
    NODE* tree;
    if (argc>1 && strcmp(argv[1],"-d")==0) debug = 1;
    init_symbtable();
    printf("--C COMPILER\n");
    yyparse();
    tree = ans;
    printf("parse finished\n");
    print_tree(tree);
    interpret_tree(tree);
    return 0;
}