#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, ...);

/*
 * 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
int 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 add_args_to_env(NODE *tree, ENV *env_ptr) {
  if (tree->type == '~') {
    BIND* arg;
    TOKEN *tok, *name;
    // we found an argument
    name = (TOKEN*)tree->right->left;
    tok = new_token(INT);
    tok->value = 0;
    asprintf(&tok->lexeme, "%d", tok->value);
    arg = create_binding(name, (NODE*) tok, NULL);
    if (env_ptr->bindings == NULL) {
      env_ptr->bindings = create_list(arg, NULL);
    } else {
      append_list(env_ptr->bindings, arg);
    }
  } else {
    add_args_to_env(tree->left, env_ptr);
    add_args_to_env(tree->right, env_ptr);
  }
}

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) {
    ENV* new_func_env = create_new_function_env(env_ptr);
    if (tree->left->right->right != NULL) {
      add_args_to_env(tree->left->right->right, new_func_env);
    }
    if (env_ptr->bindings == NULL) {
      env_ptr->bindings = create_list(create_binding(name_token, tree, new_func_env), NULL);
    } else {
      append_list(env_ptr->bindings, create_binding(name_token, tree, new_func_env));
    }
  } else {
    printf("Error: redefinition of function with name %s\n", name_token->lexeme);
  }
}

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

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

void bind_arg(TOKEN* val, TOKEN* arg_name, ENV *env_ptr) {
  BIND* existing = find_name_in_list(arg_name, env_ptr->bindings);
  if (existing == NULL) {
    printf("Attempted to bind arg that doesn't exist, the compiler is broken\n");
    exit(1);
  }
  existing->tree = (NODE*) val;
}

void add_var_to_env(NODE *tree, ENV *env_ptr) {
  BIND* existing;
  NODE* var_name = tree->left->left;
  TOKEN* name_token = (TOKEN*) var_name;
  TOKEN* tok = new_token(CONSTANT);
  tok->value = recursive_interpret(tree->right, env_ptr);
  asprintf(&tok->lexeme, "%d", tok->value);
  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*) tok, NULL), NULL);
    } else {
      append_list(env_ptr->bindings, create_binding(name_token, (NODE*) tok, NULL));
    }
  } else {
    existing->tree = (NODE *) tok;
  } 
}

int recursive_interpret(NODE *tree, ENV *env_ptr) {
  if (tree==NULL) return 0;
  if (tree->type==LEAF) {
    if (tree->left->type == CONSTANT) {
      return ((TOKEN *) tree->left)->value;
    } 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 0;
    } else {
      // an identifier
      TOKEN* tok = (TOKEN *) tree->left;
      BIND* var_bind = find_name_in_env(tok, env_ptr);
      if (var_bind == NULL) {
        printf("Could not find variable %s\n", tok->lexeme);
        exit(1);
      }
      if (var_bind->tree->type == CONSTANT) {
        TOKEN* var_tok = (TOKEN *) var_bind->tree;
        return var_tok->value;
      } else {
        printf("Maybe got a function?\n");
        return 0;
      }
    }
  }
  if (tree->type=='D') {
    // this is a function definition
    add_function_to_env(tree, env_ptr);
    return 0;
  }
  if (tree->type=='=') {
    // this is a variable definition
    add_var_to_env(tree, env_ptr);
    return 0;
  }
  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);
    }
    if (tree->right != NULL) {
      TOKEN **arg_list, **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);
      if (exp_args == num_args) {
        for (int i = 0; i < num_args; i++) {
          bind_arg(val_list[i], arg_list[i], func->env);
        }
      } else {
        printf("Incorrect arguments passed to function %s, expected %d got %d\n", func_name->lexeme, exp_args, num_args);
        exit(1);
      }
    }
    return recursive_interpret(func->tree->right, func->env);
  }
  if (tree->type == IF) {
    if (recursive_interpret(tree->left, env_ptr) == 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) {
    return (int)recursive_interpret(tree->left, env_ptr) <= recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == GE_OP) {
    return (int)recursive_interpret(tree->left, env_ptr) >= recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == EQ_OP) {
    return (int)recursive_interpret(tree->left, env_ptr) == recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == NE_OP) {
    return (int)recursive_interpret(tree->left, env_ptr) != recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '>') {
    return (int)recursive_interpret(tree->left, env_ptr) > recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '<') {
    return (int)recursive_interpret(tree->left, env_ptr) < recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '+') {
    return recursive_interpret(tree->left, env_ptr) + recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '-') {
    return recursive_interpret(tree->left, env_ptr) - recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '*') {
    return recursive_interpret(tree->left, env_ptr) * recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '/') {
    return recursive_interpret(tree->left, env_ptr) / recursive_interpret(tree->right, env_ptr);
  }
  if (tree->type == '%') {
    return recursive_interpret(tree->left, env_ptr) % recursive_interpret(tree->right, env_ptr);
  }
  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) {
    return recursive_interpret(tree->left, env_ptr);
  }
  if (tree->type == WHILE) {
    ENV* scope_env = create_new_function_env(env_ptr);
    while (recursive_interpret(tree->left, env_ptr)) {
      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);
  printf("%d\n", recursive_interpret(ref_main->tree->right, ref_main->env));
}

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) yydebug = 1;
    init_symbtable();
    printf("--C COMPILER\n");
    yyparse();
    tree = ans;
    printf("parse finished\n");
    print_tree(tree);
    interpret_tree(tree);
    return 0;
}