#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"

int debug = 0;

extern TOKEN* lookup_token(char*);

TOKEN* gen_tmp(void) {
  char tmp[10];
  static char num = 0;
  snprintf(tmp, 10, "$t%d", num++);
  return lookup_token(tmp);
}

TOKEN* gen_label(void) {
  char tmp[10];
  static char label = 0;
  snprintf(tmp, 10, "L%d", label++);
  return lookup_token(tmp);
}

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));
      print_tree0(tree->left, level+2);
      print_tree0(tree->right, level+2);
    }
}

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

void new_tac(TACS* tacs, TOKEN* dst, TOKEN* src, TOKEN* tgt, int op) {
  if (tacs->list == NULL) {
    TAC_T* tac;
    tac = create_tac(op, src, tgt, dst);
    tacs->list = new_tac_list(tac, NULL);
  } else {
    add_tac(tacs->list, op, src, tgt, dst);
  }
}

char* stok(TOKEN* tok) {
  return tok->lexeme;
}

void print_tac_el(TAC_T* elem) {
    int op = elem->op;
    if (op == '=') printf("%s := %s\n", stok(elem->dst), stok(elem->src));
    if (op=='+' || op=='-' || op=='*' || op=='/' || op =='%') printf("%s := %s %c %s\n", stok(elem->dst), stok(elem->src), op, stok(elem->tgt));
    if (op == RETURN) printf("ret %s\n", stok(elem->src));
    if (op == 'D') printf("func %s\n", stok(elem->dst));
    if (op == '>' || op == '<') printf("if %s %c %s then %s\n", stok(elem->src), op, stok(elem->tgt), stok(elem->dst));
    if (op == GE_OP) printf("if %s >= %s then %s\n", stok(elem->src), stok(elem->tgt), stok(elem->dst));
    if (op == LE_OP) printf("if %s <= %s then %s\n", stok(elem->src), stok(elem->tgt), stok(elem->dst));
    if (op == EQ_OP) printf("if %s == %s then %s\n", stok(elem->src), stok(elem->tgt), stok(elem->dst));
    if (op == NE_OP) printf("if %s != %s then %s\n", stok(elem->src), stok(elem->tgt), stok(elem->dst));
    if (op == IDENTIFIER) printf("%s\n", stok(elem->dst));
    if (op == VOID) printf("goto %s\n", stok(elem->dst));
    if (op == AUTO) printf("param %s\n", stok(elem->dst));
    if (op == APPLY) printf("call %s, %s\n", stok(elem->dst), stok(elem->tgt));
    if (op == 278) printf("arg %s\n", stok(elem->dst));
    if (op == 279) printf("end %s\n\n", stok(elem->dst));
}

void print_tac(TLIST* tac_list) {
  TLIST *ptr = tac_list;

  while (ptr != NULL) {
    print_tac_el(ptr->elem);
    ptr = ptr->next;
  }
}

int invert_op(int op) {
  if (op == LE_OP) {
    return '>';
  } else if (op == GE_OP) {
    return '<';
  } else if (op == '>') {
    return LE_OP;
  } else if (op == '<') {
    return GE_OP;
  } else if (op == EQ_OP) {
    return NE_OP;
  } else if (op == NE_OP) {
    return EQ_OP;
  } else {
    printf("Unknown operator: %c", op);
    exit(1);
  }
}

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 populate_val_list(NODE* tree, NODE** arr, int num_args) {
  if (tree->type == ',') {
    populate_val_list(tree->left, arr, num_args);
    populate_val_list(tree->right, arr, num_args);
  } else {
    for (int i = 0; i < num_args; i++) {
      if (arr[i] == NULL) {
        arr[i] = tree;
        break;
      }
    }
  }
}

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

TOKEN* build_tac(NODE *tree, TACS* tac) {
  if (tree == NULL) return NULL;
  int type = tree->type;
  if (type == LEAF) {
    return (TOKEN *) tree->left;
  } else if (type == '=') {
    TOKEN* dst = (TOKEN*) tree->left->left;
    TOKEN* result = build_tac(tree->right, tac);
    new_tac(tac, dst, result, NULL, tree->type);
  } else if (type=='+' || type=='-' || type=='*' || type=='/' || type =='%') {
    TOKEN* dst = gen_tmp();
    TOKEN* src = build_tac(tree->left, tac);
    TOKEN* tgt = build_tac(tree->right, tac);
    new_tac(tac, dst, src, tgt, tree->type);
    return dst;
  } else if (type == IF) {
    int op = invert_op(tree->left->type);
    TOKEN* lhs = build_tac(tree->left->left, tac);
    TOKEN* rhs = build_tac(tree->left->right, tac);
    TOKEN* label = gen_label();
    new_tac(tac, label, lhs, rhs, op);
    if (tree->right->type == ELSE) {
      TOKEN* cont = gen_label();
      build_tac(tree->right->left, tac);
      new_tac(tac, cont, NULL, NULL, VOID);
      new_tac(tac, label, NULL, NULL, IDENTIFIER);
      build_tac(tree->right->right, tac);
      new_tac(tac, cont, NULL, NULL, IDENTIFIER);
    } else {
      build_tac(tree->right, tac);
      new_tac(tac, label, NULL, NULL, IDENTIFIER);
    }
  } else if (type == RETURN) {
    TOKEN* src = build_tac(tree->left, tac);
    new_tac(tac, NULL, src, NULL, RETURN);
  } else if (type == 'D') {
    new_tac(tac, (TOKEN*) tree->left->right->left->left, NULL, NULL, 'D');
    if (tree->left->right->right != NULL) {
      TOKEN** arg_list = get_argument_list(tree);
      int arg_count = count_args(tree->left->right->right);
      for (int i = 0; i < arg_count; i++) {
        new_tac(tac, arg_list[i], NULL, NULL, AUTO);
      }
    }
    build_tac(tree->right, tac);
    new_tac(tac, (TOKEN*) tree->left->right->left->left, NULL, NULL, 279);
  } else if (type == APPLY) {
    TOKEN* ret_val = gen_tmp();
    if (tree->right != NULL) {
      int num_args;
      NODE **val_list;
      num_args = count_args(tree->right);
      val_list = get_val_list(tree->right);
      TOKEN* arg_tokens[num_args];
      for (int i = 0; i < num_args; i++) {
        arg_tokens[i] = build_tac(val_list[i], tac);
      }
      for (int i = 0; i < num_args; i++) {
        new_tac(tac, arg_tokens[i], NULL, NULL, 278);
      }
    }
    new_tac(tac, (TOKEN*) tree->left->left, NULL, ret_val, APPLY);
    return ret_val;
  } else {
    build_tac(tree->left, tac);
    build_tac(tree->right, tac);
    return NULL;
  }
  return NULL;
}


int tmp_regs[10] = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25};
int svd_regs[8]  = {16, 17, 18, 19, 20, 21, 22, 23};
int arg_regs[3] = {5, 6, 7};
int tmp_cntr = 0;

int register_me(TOKEN* name, ENV* env) {
  if (tmp_cntr == 10) tmp_cntr = 0;
  if (name->type == CONSTANT) {
    if (name->value == 0) {
      return 0;
    } else {
      printf("li $%d, %d\n", tmp_regs[tmp_cntr], name->value);
      return tmp_regs[tmp_cntr++];
    }
  } else {
    BIND* binding = find_name_in_list(name, env->bindings);
    if (binding == NULL) {
      BIND* new_bind;
      if (name->lexeme[0] != '$') {
        new_bind = create_binding(name, env->counter++, NULL);
        if (env->bindings == NULL) env->bindings = create_list(new_bind, NULL);
        else append_list(env->bindings, new_bind);
        new_bind->is_saved = 1;
        return svd_regs[new_bind->env_offset];
      } else {
        new_bind = create_binding(name, tmp_cntr++, NULL);
        if (env->bindings == NULL) env->bindings = create_list(new_bind, NULL);
        else append_list(env->bindings, new_bind);
        return tmp_regs[new_bind->env_offset];
      }
    } else {
      return binding->is_saved ? svd_regs[binding->env_offset] : tmp_regs[binding->env_offset];
    }
  }
}

void alloc_function_env() {
  printf("li $a0, 68\n");
  printf("li $v0, 9\n");
  printf("syscall\n");
}

void compile_tac(TLIST* tacs) {
  ENV* env = create_new_function_env(NULL, NULL);
  printf("=====\n");
  alloc_function_env();
  printf("or $fp, $0, $v0\n");
  printf("la $t0, main\n");
  printf("sw $t0, 0($fp)\n");
  printf("sw $0, 4($fp)\n");
  printf("jal main\n");
  printf("j done\n");
  while(tacs != NULL) {
    TAC_T* t = tacs->elem;
    if (debug) {
      printf("#");
      print_tac_el(t);
    }
    if (t->op == '=') {
      int src = register_me(t->src, env);
      int dst = register_me(t->dst, env);
      printf("or $%d, $0, $%d\n", dst, src);
    }
    if (t->op == 'D') {
      env = create_new_function_env(NULL, env);
      printf("%s:\n", t->dst->lexeme);
      alloc_function_env();
      printf("sw $fp, 4($v0)\n");
      if (env->previous->bindings != NULL) {
        int i = 2;
        BLIST* ptr = env->previous->bindings;
        while (ptr != NULL) {
          if (ptr->binding->name->lexeme[0] != '$') {
            printf("sw $%d, %d($fp)\n", svd_regs[ptr->binding->env_offset], i*4);
            i++;
          }
          ptr = ptr->next;
        }
      }
      printf("or $fp, $0, $v0\n");
      printf("la $t0, %s\n", stok(t->dst));
      printf("sw $t0, 0($fp)\n");
      printf("sw $ra, 64($fp)\n");
    }
    if (t->op == '+') {
      int src = register_me(t->src, env);
      int dst = register_me(t->dst, env);
      int tgt = register_me(t->tgt, env);
      printf("add $%d, $%d, $%d\n", dst, src, tgt);
    }
    if (t->op == '-') {
      int src = register_me(t->src, env);
      int dst = register_me(t->dst, env);
      int tgt = register_me(t->tgt, env);
      printf("sub $%d, $%d, $%d\n", dst, src, tgt);
    }
    if (t->op == '*') {
      int src = register_me(t->src, env);
      int dst = register_me(t->dst, env);
      int tgt = register_me(t->tgt, env);
      printf("mult $%d, $%d\n", src, tgt);
      printf("mflo $%d\n", dst);
    }
    if (t->op == '/') {
      int src = register_me(t->src, env);
      int dst = register_me(t->dst, env);
      int tgt = register_me(t->tgt, env);
      printf("div $%d, $%d\n", src, tgt);
      printf("mflo $%d\n", dst);
    }
    if (t->op == '%') {
      int src = register_me(t->src, env);
      int dst = register_me(t->dst, env);
      int tgt = register_me(t->tgt, env);
      printf("div $%d, $%d\n", src, tgt);
      printf("mfhi $%d\n", dst);
    }
    if (t->op == LE_OP) {
      int src = register_me(t->src, env);
      int tgt = register_me(t->tgt, env);
      printf("ble $%d, $%d, %s\n", src, tgt, stok(t->dst));
    }
    if (t->op == GE_OP) {
      int src = register_me(t->src, env);
      int tgt = register_me(t->tgt, env);
      printf("bge $%d, $%d, %s\n", src, tgt, stok(t->dst));
    }
    if (t->op == EQ_OP) {
      int src = register_me(t->src, env);
      int tgt = register_me(t->tgt, env);
      printf("beq $%d, $%d, %s\n", src, tgt, stok(t->dst));
    }
    if (t->op == NE_OP) {
      int src = register_me(t->src, env);
      int tgt = register_me(t->tgt, env);
      printf("bne $%d, $%d, %s\n", src, tgt, stok(t->dst));
    }
    if (t->op == '>') {
      int src = register_me(t->src, env);
      int tgt = register_me(t->tgt, env);
      printf("bgt $%d, $%d, %s\n", src, tgt, stok(t->dst));
    }
    if (t->op == '<') {
      int src = register_me(t->src, env);
      int tgt = register_me(t->tgt, env);
      printf("blt $%d, $%d, %s\n", src, tgt, stok(t->dst));
    }
    if (t->op == APPLY) {
      int tgt_idx = register_me(t->tgt, env);
      env->arg_counter = 0;
      printf("jal %s\n", stok(t->dst));
      printf("lw $fp, 4($fp)\n");
      printf("lw $ra, 64($fp)\n");
      if (env->bindings != NULL) {
        int i = 2;
        BLIST* ptr = env->bindings;
        while (ptr != NULL) {
          if (ptr->binding->name->lexeme[0] != '$') {
            printf("lw $%d, %d($fp)\n", svd_regs[ptr->binding->env_offset], i*4);
            i++;
          }
          ptr = ptr->next;
        }
      }
      printf("or $%d, $0, $v1\n", tgt_idx);
    }
    if (t->op == IDENTIFIER) {
      printf("%s:\n", stok(t->dst));
    }
    if (t->op == VOID) {
      printf("j %s\n", stok(t->dst));
    }
    if (t->op == RETURN) {
      int src_idx = register_me(t->src, env);
      printf("or $v1, $0, $%d\n", src_idx);
      printf("jr $31\n");
    }
    if (t->op == 279) {
      env = env->previous;
    }
    if (t->op == 278) {
      int dst_idx = register_me(t->dst, env);
      printf("or $%d, $0, $%d\n", arg_regs[env->arg_counter], dst_idx);
      env->arg_counter++;
    }
    if (t->op == AUTO) {
      int dst_idx = register_me(t->dst, env);
      printf("or $%d, $0, $%d\n", dst_idx, arg_regs[env->param_counter]);
      env->param_counter++;
    }
    tacs = tacs->next;
  }
  printf("done:\n");
  printf("ori $v0, $0, 10\n");
  printf("syscall\n");
}

void interpret_tree(NODE *tree) {
  TACS* gen_tac = (TACS*) malloc(sizeof(TACS));
  build_tac(tree, gen_tac);
  print_tac(gen_tac->list);
  compile_tac(gen_tac->list);
}

extern int yydebug;
#ifdef __APPLE__
extern NODE* yyparse(void);
#endif
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;
}