#include "compile.h"

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

const int BASE = 65;
const int BYTE = 8;
int IFCL_COUNT = 0;
int ELFCL_COUNT = 0;
int ENDIF_COUNT = 0;
int START_WHILE_COUNT = 0;
int END_WHILE_COUNT = 0;
int START_WHILE_BODY_COUNT = 0;
bool compile_ast(node_t *node, char *string) {
    /* You should remove this cast to void in your solution.
     * It is just here so the code compiles without warnings. */
    if (node->type == NUM) {
        printf("movq $%ld, %%rdi\n", ((num_node_t *) node)->value);
    }
    if (node->type == PRINT) {
        compile_ast(((print_node_t *) node)->expr, "PRINT");
        printf("callq print_int\n");
    }
    if (node->type == SEQUENCE) {
        for (size_t i = 0; i < ((sequence_node_t *) node)->statement_count; i++) {
            compile_ast(((sequence_node_t *) node)->statements[i], "SEQ");
        }
    }
    if (node->type == BINARY_OP) {
        compile_ast(((binary_node_t *) node)->left, "LEFT");
        // push left side
        if (((binary_node_t *) node)->right->type == NUM &&
            ((binary_node_t *) node)->op == '*') {
            node_t *right = ((binary_node_t *) node)->right;
            // int32(int) vs int64
            int64_t val = ((num_node_t *) right)->value;
            if (((val & (val - 1)) == 0) && val > 1) {
                printf("shlq $%d, %%rdi\n",
                       optimimzed_mul(((num_node_t *) right)->value));
                return true;
            }
        }
        printf("pushq %%rdi\n");
        compile_ast(((binary_node_t *) node)->right, "RIGHT");
        printf("movq %%rdi, %%rax\n");
        printf("popq %%rdi\n");
        printf("cmpq %%rax, %%rdi\n");
        // move right side to rax
        if (((binary_node_t *) node)->op == '+') {
            printf("addq %%rax, %%rdi\n");
        }
        if (((binary_node_t *) node)->op == '-') {
            printf("subq %%rax, %%rdi\n");
        }
        if (((binary_node_t *) node)->op == '*') {
            printf("imulq %%rax, %%rdi\n");
        }
        if (((binary_node_t *) node)->op == '/') {
            // move rdi to rcx
            printf("movq %%rax, %%rcx\n");
            printf("movq %%rdi, %%rax\n");
            printf("cqto\n");
            printf("idivq %%rcx\n");
            printf("movq %%rax, %%rdi\n");
        }
        if (((binary_node_t *) node)->op == '<' || ((binary_node_t *) node)->op == '=' ||
            ((binary_node_t *) node)->op == '>') {
            redirect(((binary_node_t *) node), string);
        }
    }
    if (node->type == LET) {
        compile_ast(((let_node_t *) node)->value, "LET");
        var_name_t name = ((let_node_t *) node)->var;
        if (name == 'A') {
            int offset = (name - BASE) + BYTE;
            printf("movq %%rdi, -%d(%%rbp)\n", offset);
        }
        else {
            int offset = ((name - BASE) * BYTE) + BYTE;
            printf("movq %%rdi, -%d(%%rbp)\n", offset);
        }
    }
    if (node->type == VAR) {
        var_name_t name = ((var_node_t *) node)->name;
        if (name == 'A') {
            int offset = (name - BASE) + BYTE;
            printf("movq -%d(%%rbp), %%rdi\n", offset);
        }
        else {
            int offset = ((name - BASE) * BYTE) + BYTE;
            printf("movq -%d(%%rbp), %%rdi\n", offset);
        }
    }
    if (node->type == IF) {
        binary_node_t *condition = ((if_node_t *) node)->condition;
        int copy_if = IFCL_COUNT += 1;
        int copy_else = ENDIF_COUNT += 1;
        int copy_end = ELFCL_COUNT += 1;
        compile_ast((node_t *) condition, "COND");
        // everything still needs to compile. we just need to deal with jumps somehwere.
        printf("IF_CLAUSE_%d:\n", copy_if);
        compile_ast(((if_node_t *) node)->if_branch, "IF");
        printf("jmp END_IF_CLAUSE_%d\n", copy_end);
        printf("ELSE_CLAUSE_%d:\n", copy_else);
        if (((if_node_t *) node)->else_branch) {
            compile_ast(((if_node_t *) node)->else_branch, "ELSE");
        }
        printf("END_IF_CLAUSE_%d:\n", copy_end);
    }
    if (node->type == WHILE) {
        int copy_start = START_WHILE_COUNT += 1;
        int copy_end = END_WHILE_COUNT += 1;
        int copy_body = START_WHILE_BODY_COUNT += 1;
        binary_node_t *condition = ((while_node_t *) node)->condition;
        printf("START_WHILE_%d:\n", copy_start);
        compile_ast((node_t *) condition, "WHILE");
        printf("START_WHILE_BODY_%d:\n", copy_body);
        compile_ast(((while_node_t *) node)->body, "BODY");
        printf("jmp START_WHILE_%d\n", copy_start);
        printf("END_WHILE_%d:\n", copy_end);
    }
    return true; // for now, every statement causes a compilation failure
}

void redirect(binary_node_t *node, char *string) {
    if (((binary_node_t *) node)->op == '<') {
        if (strcmp(string, "WHILE") == 0) {
            printf("jl START_WHILE_BODY_%d\n", START_WHILE_BODY_COUNT);
            printf("jmp END_WHILE_%d\n", END_WHILE_COUNT);
        }
        else {
            printf("jl IF_CLAUSE_%d\n", IFCL_COUNT);
            printf("jmp ELSE_CLAUSE_%d\n", ELFCL_COUNT);
        }
    }
    if (((binary_node_t *) node)->op == '=') {
        if (strcmp(string, "WHILE") == 0) {
            printf("je START_WHILE_BODY_%d\n", START_WHILE_BODY_COUNT);
            printf("jmp END_WHILE_%d\n", END_WHILE_COUNT);
        }
        else {
            printf("je IF_CLAUSE_%d\n", IFCL_COUNT);
            printf("jmp ELSE_CLAUSE_%d\n", ELFCL_COUNT);
        }
    }
    if (((binary_node_t *) node)->op == '>') {
        if (strcmp(string, "WHILE") == 0) {
            printf("jg START_WHILE_BODY_%d\n", START_WHILE_BODY_COUNT);
            printf("jmp END_WHILE_%d\n", END_WHILE_COUNT);
        }
        else {
            printf("jg IF_CLAUSE_%d\n", IFCL_COUNT);
            printf("jmp ELSE_CLAUSE_%d\n", ELFCL_COUNT);
        }
    }
}

int optimimzed_mul(int64_t x) {
    int count = 0;
    while ((x > 1) && (x % 2 == 0)) {
        count++;
        x = x >> 1;
    }
    return count;
}

// void traverser(node_t *tnode) {
//     int result = 0;
//     traverser_helper(tnode, result);
// }

// void traverser_helper(node_t *node, int result) {
//     if (node->type == NUM){
//         result += ((num_node_t*)node)->value;
//     }
//     if (((binary_node_t*)node)->left->type != NUM && ((binary_node_t*)node)->right->type != NUM) {
//         traverser_helper(((binary_node_t*)node)->left, result);
//         traverser_helper(((binary_node_t*)node)->right, result);
//     }
//     if (node->type == BINARY_OP){
//         binary_node_t* node = node;
//         switch (((binary_node_t*)node)->op)
//         {
//         case '+':
//             // would we init num node here?
//             init_num_node(result);
//             break;
//         case '-':
//             init_num_node(result);
//             break;
//         case '*':
//             init_num_node(result);
//             break;
//         case '/':
//             init_num_node(result);
//             break;
//         default:
//             break;
//         }
//     }
// }