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

#include <mcc/ast.h>
#include "mcc/tac_build.h"
#include "mcc/tac_string_builder.h"
#include "mcc/symbol_table_validate.h"
#include "mcc/symbol_table_print.h"
#include "mcc/utils.h"

/*
 * Source for multiple parts:
 * https://web.stanford.edu/class/archive/cs/cs143/cs143.1128/lectures/13/Slides13.pdf
 */

// -------------------------------- parse expression

void mcc_tac_parse_expression(struct mcc_ast_expression *expression, struct mcc_tac *tac) {
    assert(expression);
    assert(tac);

    switch(expression->type) {
        case MCC_AST_EXPRESSION_TYPE_IDENTIFIER: {
            struct mcc_ast_identifier *identifier = expression->identifier;

            enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(expression,
                                                                                      tac->current_symbol_table);

            enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_ident(type);

            char *arg1 = identifier->i_value;
            mcc_tac_create_and_add_new_entry_temp(op, arg1, NULL, tac);

        } break;
        case MCC_AST_EXPRESSION_TYPE_LITERAL : {
            struct mcc_ast_literal *literal = expression->literal;
            enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_literal(literal->type);
            switch(op) {
                case MCC_TAC_BOOL :
                    mcc_tac_create_and_add_new_entry_temp(op, literal -> b_value ? "true" : "false", NULL, tac);
                    break;
                case MCC_TAC_INT : {
                    // + 2 for \0 and in case result is 0
                    char arg1[mcc_get_number_of_digits(literal -> i_value) + 2];
                    sprintf(arg1, "%d", literal -> i_value);

                    mcc_tac_create_and_add_new_entry_temp(op, arg1, NULL, tac);
                } break;
                case MCC_TAC_FLOAT : {
                    // just for conversion - when storing the value in tac_new_entry it will take actual string length
                    char arg1[1000];
                    gcvt(literal->f_value, 6, arg1);
                    mcc_tac_create_and_add_new_entry_temp(op, arg1, NULL, tac);
                } break;
                case MCC_TAC_STRING :
                    mcc_tac_create_and_add_new_entry_temp(op, literal -> s_value, NULL, tac);
                    break;
                default:
                    break;
            }
        } break;
        case MCC_AST_EXPRESSION_TYPE_CALL_EXPRESSION : {
           // check if call expression has arguments and push/pop parameters
            int arg_size = 0;
            if(expression->argument != NULL) {
                struct mcc_ast_argument *argument = expression -> argument;
                arg_size = argument -> expressions -> size;

                // push args
                for(int i = arg_size - 1; i >= 0; --i) {
                    struct mcc_ast_expression *parameter = argument->expressions->arr[i];
                    bool is_array = parameter -> bracket_expression != NULL;

                    enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(parameter,
                                                                                              tac->current_symbol_table);

                    enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_param_push(type, is_array);

                    mcc_tac_parse_expression(parameter, tac);
                    char *arg1 = tac->last_temporary;
                    mcc_tac_create_and_add_new_entry(op, arg1, NULL, NULL, tac);
                }
            }

            {
                char *result = expression->function_name->i_value;

                mcc_tac_create_and_add_new_entry(MCC_TAC_CALL, NULL, NULL, result, tac);

                enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(expression,
                                                                                          tac->current_symbol_table);
                enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_ident(type);

                if(op != MCC_TAC_UNKNOWN) {
                    char *arg1 = mcc_tac_new_return_function_name(result);
                    mcc_tac_create_and_add_new_entry_temp(op, arg1, NULL, tac);
                }
            }

            // POP params
            if(expression->argument != NULL) {
                struct mcc_ast_argument *argument = expression -> argument;
                arg_size = argument -> expressions -> size;

                for(int i = arg_size - 1; i >= 0; --i) {
                    struct mcc_ast_expression *parameter = argument->expressions->arr[i];
                    bool is_array = parameter -> bracket_expression != NULL;

                    enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(parameter,
                                                                                              tac->current_symbol_table);
                    enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_param_pop(type, true);

                    // get array size from symbol table
                    long array_size = 0;
                    if (is_array) {
                        struct mcc_symbol *s = mcc_symbol_table_get_symbol(tac -> current_symbol_table,
                                                                           parameter -> bracket_identifier ->i_value, is_array);

                        if (s != NULL) {
                            array_size = s ->array_size;
                        }

                        char arg1[mcc_get_number_of_digits((int) array_size) + 2];
                        sprintf(arg1, "%ld", array_size);
                        mcc_tac_create_and_add_new_entry(op, arg1, NULL, NULL, tac);
                    } else {
                        // printf("%d\n", op);
                        mcc_tac_create_and_add_new_entry(op, NULL, NULL, NULL, tac);
                    }
                }
            }

        } break;
        case MCC_AST_EXPRESSION_TYPE_UNARY_OP : {

            enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(expression,
                                                                                      tac->current_symbol_table);
            enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_unary(type);

            mcc_tac_parse_expression(expression->unary_expression, tac);
            char *arg1 = tac->last_temporary;
            mcc_tac_create_and_add_new_entry_temp(op, arg1, NULL, tac);
        }
            break;
        case MCC_AST_EXPRESSION_TYPE_BINARY_OP : {

            enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(expression,
                                                                                      tac->current_symbol_table);
            enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_binary(expression->op, type);

            mcc_tac_parse_expression(expression->lhs, tac);
            char *arg1 = tac->last_temporary;

            mcc_tac_parse_expression(expression->rhs, tac);
            char *arg2 = tac->last_temporary;

            mcc_tac_create_and_add_new_entry_temp(op, arg1, arg2, tac);

        }
            break;
        case MCC_AST_EXPRESSION_TYPE_PARENTH :
            mcc_tac_parse_expression(expression->expression, tac);
            break;
        case MCC_AST_EXPRESSION_TYPE_BRACKET : {
            enum mcc_ast_data_type type = mcc_symbol_table_get_expression_return_type(expression,
                                                                                      tac->current_symbol_table);
            enum mcc_tac_operation op = mcc_convert_ast_type_to_tac_load(type);

            struct mcc_ast_identifier *ident = expression->bracket_identifier;
            char *arg1 = ident->i_value;

            mcc_tac_parse_expression(expression->bracket_expression, tac);

            char *arg2 = tac->last_temporary;

            mcc_tac_create_and_add_new_entry_temp(op, arg1, arg2, tac);
        }
            break;
        default:
            return;
    }
}

// -------------------------------- parse statement

void mcc_tac_parse_statement_list(struct mcc_ast_statement_list *stl, struct mcc_tac *tac) {
    assert(tac);

    while(stl != NULL) {
        mcc_tac_parse_statement(stl->statement, tac);
        stl = stl->next;
    }
}

void mcc_tac_parse_statement(struct mcc_ast_statement *statement, struct mcc_tac *tac) {
    assert(tac);

    switch(statement->type) {
        case MCC_AST_STATEMENT_TYPE_COMPOUND:
            mcc_tac_parse_statement_list(statement->statement_list, tac);
            break;
        case MCC_AST_STATEMENT_TYPE_IF: {
            // evaluate if expression
            mcc_tac_parse_expression(statement->if_condition, tac);

            char *last_temp = tac->last_temporary;

            // end of if statement code
            char *if_label = mcc_tac_new_label_string(tac);

            mcc_tac_create_and_add_new_entry(MCC_TAC_JMP_FALSE, last_temp, NULL, if_label, tac);

            // scope go to deeper scope and save current index
            int current_scope_index = tac->current_symbol_index;

            mcc_tac_enter_next_deeper_scope(tac);

            mcc_tac_parse_statement(statement->if_stmt, tac);
            // go back to previous scope
            mcc_tac_exit_to_outer_scope(tac, current_scope_index);

            tac->current_symbol_index++;

            char *else_label = mcc_tac_new_label_string(tac);
            mcc_tac_create_and_add_new_entry(MCC_TAC_JMP, NULL, NULL, else_label, tac);

            // end if code with if_label
            mcc_tac_create_and_add_new_entry(MCC_TAC_LABEL, NULL, NULL, if_label, tac);

            if(statement->else_stmt != NULL) {
                current_scope_index = tac->current_symbol_index;

                mcc_tac_enter_next_deeper_scope(tac);

                mcc_tac_parse_statement(statement->else_stmt, tac);

                //go back to previous scope
                mcc_tac_exit_to_outer_scope(tac, current_scope_index);
                tac->current_symbol_index++;
            }

            // set else label anyway
            mcc_tac_create_and_add_new_entry(MCC_TAC_LABEL, NULL, NULL, else_label, tac);

        }
            break;
        case MCC_AST_STATEMENT_TYPE_EXPRESSION:
            mcc_tac_parse_expression(statement->expression, tac);

            break;
        case MCC_AST_STATEMENT_TYPE_WHILE: {
            char *before_while_label = mcc_tac_new_label_string(tac);
            char *after_while_label = mcc_tac_new_label_string(tac);

            mcc_tac_create_and_add_new_entry(MCC_TAC_LABEL, NULL, NULL, before_while_label, tac);
            mcc_tac_parse_expression(statement->while_condition, tac);

            char *last_temp = tac->last_temporary;

            mcc_tac_create_and_add_new_entry(MCC_TAC_JMP_FALSE, last_temp, NULL, after_while_label, tac);

            int current_scope_index = tac->current_symbol_index;
            mcc_tac_enter_next_deeper_scope(tac);

            mcc_tac_parse_statement(statement->while_stmt, tac);

            //go back to previous scope
            mcc_tac_exit_to_outer_scope(tac, current_scope_index);
            tac->current_symbol_index++;

            // jump back up
            mcc_tac_create_and_add_new_entry(MCC_TAC_JMP, NULL, NULL, before_while_label, tac);

            // add end while label
            mcc_tac_create_and_add_new_entry(MCC_TAC_LABEL, NULL, NULL, after_while_label, tac);

        }
            break;
        case MCC_AST_STATEMENT_TYPE_DECL: {
            struct mcc_ast_declaration *declaration = statement -> declaration;
            struct mcc_ast_literal *arr_literal = declaration->arr_literal;

            // Add entry for array size
            if(arr_literal != NULL) {
                // save array size in first argument
                char arg1[mcc_get_number_of_digits(arr_literal->i_value) + 2];
                sprintf(arg1, "%d", arr_literal->i_value);

                mcc_tac_create_and_add_new_entry(
                        MCC_TAC_ARR_DECL,
                        mcc_ast_type_to_string_size(declaration ->type),
                        NULL,
                        arg1,
                        tac
                );
            }
        }
            break;
        case MCC_AST_STATEMENT_TYPE_ASSGN_ARR:
        case MCC_AST_STATEMENT_TYPE_ASSGN: {
            // get type
            struct mcc_ast_assignment *a = statement->assignment;
            char *result = statement->assignment->identifier->i_value;

            enum mcc_ast_data_type assignment_type = mcc_symbol_table_get_symbol(
                    tac->current_symbol_table,
                    result,
                    true)->data_type;

            enum mcc_tac_operation tac_op = mcc_convert_ast_type_to_tac_assignment(assignment_type);
            if(a->type == MCC_AST_ASSIGNMENT_TYPE_ARRAY) {
                mcc_tac_parse_expression(a->array_ass.index, tac);

                char *arg2 = tac->last_temporary;

                mcc_tac_parse_expression(a->array_ass.rhs, tac);

                char *arg1 = tac->last_temporary;
                mcc_tac_create_and_add_new_entry(tac_op, arg1, arg2, result, tac);

            } else {

                // evaluate expression
                mcc_tac_parse_expression(a->normal_ass.rhs, tac);

                char *arg1 = tac->last_temporary;
                mcc_tac_create_and_add_new_entry(tac_op, arg1, NULL, result, tac);
            }

        }
            break;
        case MCC_AST_STATEMENT_TYPE_RETURN: {
            mcc_tac_parse_expression(statement->expression, tac);
            char *arg1 = tac->last_temporary;

            char *result = mcc_tac_new_return_function_name(tac->current_symbol_table->sym_table_name);
            mcc_tac_create_and_add_new_entry(MCC_TAC_RETURN, arg1, NULL, result, tac);
        } break;
    }

}

// -------------------------------- parse function

// TODO: check if not used when assembly code is generated (probably not) -> then kick this out
//void parse_params(struct mcc_ast_parameter *parameter, struct mcc_tac *tac) {
//    // parse params from right to left
//    for(int i = 0; i < parameter->parameters->size; i++) {
//        struct mcc_ast_declaration *param_decl = (struct mcc_ast_declaration *) parameter->parameters->arr[i];
//        bool is_array = param_decl->arr_literal != NULL;
//        enum mcc_tac_operation tac_op = mcc_convert_ast_type_to_tac_param(param_decl->type, is_array);
//
//        // save array size in first argument
//        int number_of_digits = mcc_get_number_of_digits(param_decl->arr_literal->i_value);
//        char *arg1 = NULL;
//
//        if(is_array) {
//            snprintf(arg1, sizeof(char) * number_of_digits, "%d", param_decl->arr_literal->i_value);
//        }
//
//        mcc_tac_create_and_add_new_entry(tac_op, arg1, NULL, param_decl->ident->i_value, tac);
//    }
//}

void mcc_tac_parse_function(struct mcc_ast_function *f, struct mcc_tac *tac) {
    assert(f);
    assert(tac);

    // set function symbol table scope
    struct mcc_symbol_table *t = mcc_symbol_table_get_inner_table_by_name(tac->root_table, f->identifier->i_value);
    tac->current_symbol_table = t;
    tac->current_symbol_index = 0;

    char arg1[mcc_get_number_of_digits(tac->funcion_counter) + 2];
    sprintf(arg1, "%d", tac -> funcion_counter);
    mcc_tac_create_and_add_new_entry(MCC_TAC_FUNCTION_START, arg1, f -> identifier -> i_value, NULL, tac);


    if(f->statement != NULL) {
        mcc_tac_parse_statement(f->statement, tac);
    }


    Dynamic_Array *tac_entries = tac -> tac_entries;
    struct mcc_tac_entry *last_tac_entry = (struct mcc_tac_entry *) tac_entries -> arr[tac_entries -> size - 1];

    // add function-end operation if last operation is not return
    if (last_tac_entry -> tac_op != MCC_TAC_RETURN) {
        mcc_tac_create_and_add_new_entry(MCC_TAC_FUNCTION_END, NULL, NULL, NULL, tac);
    }

    tac -> funcion_counter += 1;
}

// start parsing program
struct mcc_tac *mcc_tac_build(struct mcc_ast_program *program, struct mcc_symbol_table *st) {
    struct mcc_tac *tac = mcc_tac_new(NULL, st);

    tac->root_table = st;

    struct mcc_ast_function *f = NULL;

    // start with main
    for(int i = 0; i < program->function_def->size; i++) {
        f = program->function_def->arr[i];

        // parse all functions one after another
        mcc_tac_parse_function(f, tac);
    
    }
    // mcc_tac_create_and_add_new_entry(MCC_TAC_EOF,NULL,NULL,NULL,tac);

    return tac;
}