#include "../include/program.h"
#include "../include/bytecode.h"
#include "../include/command.h"
#include "../include/std_gc/list.h"
#ifdef DEBUG
#define GC_DEBUG
#endif
#include <gc.h>
#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>

#define cleanup_if_error(expr)                  \
    {                                           \
        if(!(expr)) {                           \
            result = TSL_PROGRAM_RESULT_ERR;    \
            goto cleanup;                       \
        }                                       \
    }

#define return_if_error(expr)                   \
    {                                           \
        if(!(expr))                             \
            return TSL_PROGRAM_RESULT_ERR;      \
    }

void tsl_program_init(TslProgram *self) {
    GC_INIT(); /*  TODO: Remove this */
    tsl_buffer_init(&self->function_bytecode_list);
}

void tsl_program_deinit(TslProgram *self) {
    TslBytecode *bytecode_writer = tsl_buffer_begin(&self->function_bytecode_list);
    TslBytecode *bytecode_writer_end = tsl_buffer_end(&self->function_bytecode_list);
    while(bytecode_writer != bytecode_writer_end) {
        tsl_bytecode_deinit(bytecode_writer);
        ++bytecode_writer;
    }
    tsl_buffer_deinit(&self->function_bytecode_list);
    GC_deinit(); /* TODO: Remove this */
}

static int program_output_temp(char *data, int size, void *userdata) {
    (void)size;
    (void)userdata;
    fputs(data, stdout);
    return 1;
}

static void tsl_string_to_ref(const TslString *self, TslStringView *result) {
    result->data = self->data;
    result->size = self->size;
}

static TslProgramResult tsl_string_ref_get_item_at_index(const TslStringView *self, const TslValue *key, TslValue *result) {
    TslValue *list_item;
    if(key->type != TSL_TYPE_NUMBER) {
        /* TODO: Print type as string */
        fprintf(stderr, "Error: Unable to index string %s using a variable of type %d. The index has to be a number\n", "TODO", key->type);
        return TSL_PROGRAM_RESULT_ERR;
    }

    list_item = tsl_string_ref_get(self, key->data.number);
    if(list_item) {
        *result = *list_item;
    } else {
        result->type = TSL_TYPE_NULL;
    }

    return TSL_PROGRAM_RESULT_OK;
}

static TslProgramResult tsl_value_create_from_stack_value(TslProgram *self, TslValue *dst) {
    TslStackValue *src = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue));
    tsl_value_clear(dst);
    switch(src->type) {
        case TSL_STACK_VALUE_TYPE_NUMBER: {
            dst->data.number = src->data.number;
            dst->type = TSL_TYPE_NUMBER;
            break;
        }
        case TSL_STACK_VALUE_TYPE_BOOL: {
            dst->data.boolean = src->data.boolean;
            dst->type = TSL_TYPE_BOOL;
            break;
        }
        case TSL_STACK_VALUE_TYPE_STRING: {
            #if 0
            TslString *str = GC_MALLOC_ATOMIC(sizeof(TslString));
            str->data = GC_MALLOC_ATOMIC(src->data.str.size + 1);
            if(!str) {
                GC_FREE(str);
                GC_FREE(str->data);
                fprintf(stderr, "Error: out of memory\n");
                return TSL_PROGRAM_RESULT_ERR;
            }
            if(!str->data) {
                GC_FREE(str);
                GC_FREE(str->data);
                fprintf(stderr, "Error: out of memory\n");
                return TSL_PROGRAM_RESULT_ERR;
            }

            memcpy(str->data, src->data.str.data, src->data.str.size);
            str->data[src->data.str.size] = '\0';
            str->size = src->data.str.size;

            dst->data.string = str;
            dst->type = TSL_TYPE_STRING;
            #endif
            dst->data.string_ref = src->data.str;
            dst->type = TSL_TYPE_STRING_REF;
            break;
        }
        case TSL_STACK_VALUE_TYPE_FUNCTION: {
            dst->data.function = src->data.integer;
            dst->type = TSL_TYPE_FUNCTION;
            break;
        }
        case TSL_STACK_VALUE_TYPE_VARIABLE_NAME: {
            TslValue *var;
            TslValue map_key;
            map_key.type = TSL_TYPE_STRING_REF;
            map_key.data.string_ref = src->data.str;

            var = tsl_hash_map_get(self->variables, &map_key);
            if(!var) {
                fprintf(stderr, "Error: Trying to access a non-existing variable \"%.*s\"\n", (int)src->data.str.size, src->data.str.data);
                return TSL_PROGRAM_RESULT_ERR;
            }

            *dst = *var;
            break;
        }
        case TSL_STACK_VALUE_TYPE_VARIABLE: {
            *dst = src->data.variable;
            break;
        }
        case TSL_STACK_VALUE_TYPE_LIST: {
            TslValue *list_item;
            int i = src->data.integer - 1;
            dst->data.list = GC_MALLOC(sizeof(TslList));
            return_if_error(dst->data.list);

            tsl_list_init(dst->data.list);
            return_if_error(tsl_list_set_capacity_hint(dst->data.list, sizeof(TslValue) * src->data.integer));
            dst->data.list->size = sizeof(TslValue) * src->data.integer;
            list_item = tsl_list_end(dst->data.list) - sizeof(TslValue);

            while(i >= 0) {
                return_if_error(tsl_value_create_from_stack_value(self, list_item));
                --list_item;
                --i;
            }

            dst->type = TSL_TYPE_LIST;
            break;
        }
        case TSL_STACK_VALUE_TYPE_MAP: {
            int i = src->data.integer - 1;

            dst->data.map = GC_MALLOC(sizeof(TslHashMap));
            return_if_error(dst->data.map);
            tsl_hash_map_init(dst->data.map);

            assert(src->data.integer % 2 == 0);
            while(i >= 0) {
                TslValue map_key;
                TslValue map_value;
                return_if_error(tsl_value_create_from_stack_value(self, &map_value));
                return_if_error(tsl_value_create_from_stack_value(self, &map_key));
                return_if_error(tsl_hash_map_insert(dst->data.map, &map_key, &map_value));
                i -= 2;
            }

            dst->type = TSL_TYPE_MAP;
            break;
        }
        case TSL_STACK_VALUE_TYPE_INDEX: {
            TslValue var;
            TslValue key;
            return_if_error(tsl_value_create_from_stack_value(self, &key));
            return_if_error(tsl_value_create_from_stack_value(self, &var));

            if(var.type == TSL_TYPE_LIST) {
                TslValue *list_item;
                if(key.type != TSL_TYPE_NUMBER) {
                    /* TODO: Print type as string */
                    fprintf(stderr, "Error: Unable to index list %s using a variable of type %d. The index has to be a number\n", "TODO", key.type);
                    return TSL_PROGRAM_RESULT_ERR;
                }

                list_item = tsl_list_get(var.data.list, key.data.number);
                if(list_item) {
                    *dst = *list_item;
                } else {
                    dst->type = TSL_TYPE_NULL;
                }
                return TSL_PROGRAM_RESULT_OK;
            } else if(var.type == TSL_TYPE_MAP) {
                TslValue *map_value = tsl_hash_map_get(var.data.map, &key);
                if(map_value) {
                    *dst = *map_value;
                } else {
                    dst->type = TSL_TYPE_NULL;
                }
                return TSL_PROGRAM_RESULT_OK;
            } else if(var.type == TSL_TYPE_STRING) {
                TslStringView str_ref;
                tsl_string_to_ref(var.data.string, &str_ref);
                return tsl_string_ref_get_item_at_index(&str_ref, &key, dst);
            } else if(var.type == TSL_TYPE_STRING_REF) {
                return tsl_string_ref_get_item_at_index(&var.data.string_ref, &key, dst);
            } else {
                /* TODO: Print type as string */
                fprintf(stderr, "Error: Unable to index data of type %d. Expected list or map\n", var.type);
                return TSL_PROGRAM_RESULT_ERR;
            }
        }
        case TSL_STACK_VALUE_TYPE_NULL: {
            dst->type = TSL_TYPE_NULL;
            break;
        }
        case TSL_STACK_VALUE_TYPE_COMMAND_ARG: {
            assert(0 && "This is a bug in the compiler. Unable to convert command arg to tsl value");
            abort();
            break;
        }
    }
    return TSL_PROGRAM_RESULT_OK;
}

static TslProgramResult tsl_program_perform_operation(TslProgram *self, TslNumber(*operation_func)(TslNumber lhs, TslNumber rhs)) {
    TslValue lhs_value;
    TslValue rhs_value;
    TslStackValue stack_value;

    return_if_error(tsl_value_create_from_stack_value(self, &rhs_value));
    return_if_error(tsl_value_create_from_stack_value(self, &lhs_value));

    if(lhs_value.type != TSL_TYPE_NUMBER) {
        fprintf(stderr, "Error: Unable to perform operation '+' between a TODO and a TODO\n");
        return TSL_PROGRAM_RESULT_ERR;
    }
    if(rhs_value.type != TSL_TYPE_NUMBER) {
        fprintf(stderr, "Error: Unable to perform operation '+' between a TODO and a TODO\n");
        return TSL_PROGRAM_RESULT_ERR;
    }

    stack_value.data.number = operation_func(lhs_value.data.number, rhs_value.data.number);
    stack_value.type = TSL_STACK_VALUE_TYPE_NUMBER;
    return tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value));
}

static TslNumber number_operation_sub(TslNumber lhs, TslNumber rhs) {
    return lhs + rhs;
}

static TslNumber number_operation_mul(TslNumber lhs, TslNumber rhs) {
    return lhs * rhs;
}

static TslNumber number_operation_div(TslNumber lhs, TslNumber rhs) {
    return lhs / rhs;
}

static TslProgramResult tsl_program_run_function(TslProgram *self, int function_index) {
    TslProgramResult result = TSL_PROGRAM_RESULT_OK;
    TslBytecode *function_bytecode = (TslBytecode*)tsl_buffer_begin(&self->function_bytecode_list) + function_index;
    char *instruction = tsl_buffer_begin(&function_bytecode->buffer);
    char *instruction_end = tsl_buffer_end(&function_bytecode->buffer);
    size_t prev_stack_value_size = self->stack_values.size;

    /* TODO: Verify if these don't cause unaligned memory access on non-x86 platforms */
    while(instruction != instruction_end) {
        TslOpcode opcode = *(TslOpcode*)instruction;
        TslInstructionType1 *instruction_type1 = (TslInstructionType1*)instruction;
        TslInstructionType2 *instruction_type2 = (TslInstructionType2*)instruction;
        TslInstructionType3 *instruction_type3 = (TslInstructionType3*)instruction;
        TslInstructionType4 *instruction_type4 = (TslInstructionType4*)instruction;
        switch(opcode) {
            case TSL_OPCODE_LOADN: {
                TslStackValue stack_value;
                stack_value.data.number = instruction_type2->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_NUMBER;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loadn %f\n", instruction_type2->value);*/
                instruction += sizeof(TslInstructionType2);
                break;
            }
            case TSL_OPCODE_LOADB: {
                TslStackValue stack_value;
                stack_value.data.boolean = instruction_type3->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_BOOL;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loadb %s\n", instruction_type3->value ? "true" : "false");*/
                instruction += sizeof(TslInstructionType3);
                break;
            }
            case TSL_OPCODE_LOADS: {
                TslStackValue stack_value;
                stack_value.data.str = instruction_type4->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_STRING;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loads \"%.*s\"\n", (int)instruction_type4->value.size, instruction_type4->value.data);*/
                instruction += sizeof(TslInstructionType4);
                break;
            }
            case TSL_OPCODE_LOADF: {
                TslStackValue stack_value;
                stack_value.data.integer = instruction_type1->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_FUNCTION;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loadf %d\n", instruction_type1->value);*/
                instruction += sizeof(TslInstructionType1);
                break;
            }
            case TSL_OPCODE_LOADV: {
                TslStackValue stack_value;
                stack_value.data.str = instruction_type4->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_VARIABLE_NAME;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loadv \"%.*s\"\n", (int)instruction_type4->value.size, instruction_type4->value.data);*/
                instruction += sizeof(TslInstructionType4);
                break;
            }
            case TSL_OPCODE_LOADNULL: {
                TslStackValue stack_value;
                stack_value.type = TSL_STACK_VALUE_TYPE_NULL;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loadnull\n");*/
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_SETV: {
                TslValue map_key;
                TslValue *map_value;
                map_key.type = TSL_TYPE_STRING_REF;
                map_key.data.string_ref = instruction_type4->value;
                
                map_value = tsl_hash_map_get_or_create(self->variables, &map_key);
                cleanup_if_error(map_value);

                cleanup_if_error(tsl_value_create_from_stack_value(self, map_value));
                /*printf("setv \"%.*s\"\n", (int)instruction_type4->value.size, instruction_type4->value.data);*/
                instruction += sizeof(TslInstructionType4);
                break;
            }
            case TSL_OPCODE_LIST: {
                TslStackValue stack_value;
                stack_value.data.integer = instruction_type1->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_LIST;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("list %d\n", instruction_type1->value);*/
                instruction += sizeof(TslInstructionType1);
                break;
            }
            case TSL_OPCODE_MAP: {
                TslStackValue stack_value;
                stack_value.data.integer = instruction_type1->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_MAP;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("map %d\n", instruction_type1->value);*/
                instruction += sizeof(TslInstructionType1);
                break;
            }
            case TSL_OPCODE_MINDEX: {
                TslStackValue stack_value;
                stack_value.type = TSL_STACK_VALUE_TYPE_INDEX;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("mindex\n");*/
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_CALLF: {
                size_t args_size = 1 + instruction_type1->value;
                TslStackValue *args = (TslStackValue*)tsl_buffer_end(&self->stack_values) - args_size;
                assert(self->stack_values.size >= args_size);

                {
                    /* Resolved values pushed to stack and replace them in the stack */
                    size_t prev_stack_size = self->stack_values.size;
                    TslStackValue *arg = (TslStackValue*)tsl_buffer_end(&self->stack_values) - 1;
                    TslStackValue *args_begin = args - 1;
                    while(arg != args_begin) {
                        TslValue new_value;
                        cleanup_if_error(tsl_value_create_from_stack_value(self, &new_value));
                        arg->data.variable = new_value;
                        arg->type = TSL_STACK_VALUE_TYPE_VARIABLE;
                        --arg;
                    }
                    /*
                        Reset the stack size here because @tsl_value_create_from_stack_value which is called above pops the stack (without reallocation).
                        TODO: Modify @tsl_value_create_from_stack_value to not pop stack and instead operate with stack pointer.
                    */
                    self->stack_values.size = prev_stack_size;
                }

                assert(args[0].type == TSL_STACK_VALUE_TYPE_VARIABLE);
                if(args[0].data.variable.type != TSL_TYPE_FUNCTION) {
                    fprintf(stderr, "Error: Unable to call a non-function type\n");
                    goto cleanup;
                }

                cleanup_if_error(tsl_program_run_function(self, args[0].data.variable.data.function));
                tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * args_size);

                /* TODO: Implement this */
                {
                    /* TODO: This is only temporary. This should be replaced with push the result of the function call onto the stack */
                    TslStackValue stack_value;
                    stack_value.type = TSL_STACK_VALUE_TYPE_NULL;
                    cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                }
                /*printf("callf %d\n", instruction_type1->value);*/
                instruction += sizeof(TslInstructionType1);
                break;
            }
            case TSL_OPCODE_ADD: {
                TslValue lhs_value;
                TslValue rhs_value;
                TslStackValue stack_value;

                cleanup_if_error(tsl_value_create_from_stack_value(self, &rhs_value));
                cleanup_if_error(tsl_value_create_from_stack_value(self, &lhs_value));
                cleanup_if_error(tsl_value_add(&lhs_value, &rhs_value, &stack_value.data.variable));

                stack_value.type = TSL_STACK_VALUE_TYPE_VARIABLE;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("add\n");*/
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_SUB: {
                cleanup_if_error(tsl_program_perform_operation(self, number_operation_sub));
                /*printf("sub\n");*/
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_MUL: {
                cleanup_if_error(tsl_program_perform_operation(self, number_operation_mul));
                /*printf("mul\n");*/
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_DIV: {
                cleanup_if_error(tsl_program_perform_operation(self, number_operation_div));
                /*printf("div\n");*/
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_LOADCA: {
                TslStackValue stack_value;
                stack_value.data.str = instruction_type4->value;
                stack_value.type = TSL_STACK_VALUE_TYPE_COMMAND_ARG;
                cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                /*printf("loadca \"%.*s\"\n", (int)instruction_type4->value.size, instruction_type4->value.data);*/
                instruction += sizeof(TslInstructionType4);
                break;
            }
            case TSL_OPCODE_CALLC: {
                TslStackValue *args_begin = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * instruction_type1->value);
                TslStackValue *args_end = args_begin + instruction_type1->value;
                TslStackValue *arg = args_begin;
                char *command_args[255];
                char temp_null_save[255];
                size_t arg_index = 0;
                assert(instruction_type1->value > 0);
                {
                    /* TODO: Support infinite amount of args */
                    assert(args_end - args_begin < 255);
                    /*printf("Running command: ");*/
                    while(arg != args_end) {
                        assert(arg->type == TSL_STACK_VALUE_TYPE_COMMAND_ARG);
                        temp_null_save[arg_index] = arg->data.str.data[arg->data.str.size];
                        arg->data.str.data[arg->data.str.size] = '\0';
                        
                        command_args[arg_index] = arg->data.str.data;
                        /*printf("%s ", command_args[arg_index]);*/
                        ++arg_index;
                        ++arg;
                    }
                    command_args[arg_index] = NULL;
                    /*printf("\n");*/
                }

                tsl_command_exec(command_args, program_output_temp, NULL);

                {
                    arg = args_begin;
                    arg_index = 0;
                    /* TODO: Support infinite amount of args */
                    while(arg != args_end) {
                        arg->data.str.data[arg->data.str.size] = temp_null_save[arg_index];
                        ++arg_index;
                        ++arg;
                    }
                    command_args[arg_index] = NULL;
                }

                {
                    /* TODO: This is only temporary. This should be replaced with push the result of the command onto the stack */
                    TslStackValue stack_value;
                    stack_value.type = TSL_STACK_VALUE_TYPE_NULL;
                    cleanup_if_error(tsl_buffer_append(&self->stack_values, &stack_value, sizeof(stack_value)));
                }
                
                /*printf("callc %d\n", instruction_type1->value);*/
                instruction += sizeof(TslInstructionType1);
                break;
            }
        }
    }

    cleanup:
    /*
        This a bug in the compiler. The compiler pop'ed more values that pushed to the stack.
        It's fine if there were values that were pushed to the stack but not pop'ed. This happens for example if there is a function call
        and the result is not assigned to a variable (in other words, the result is ignored).
        This will automatically be cleaned up here and in loops it will be cleaned up at the end of the loop scope.
    */
    assert(self->stack_values.size >= prev_stack_value_size);
    self->stack_values.size = prev_stack_value_size;
    return result;
}

TslProgramResult tsl_program_run(TslProgram *self) {
    TslProgramResult result = TSL_PROGRAM_RESULT_OK;

    self->variables = GC_MALLOC_UNCOLLECTABLE(sizeof(TslHashMap));
    if(!self->variables) {
        fprintf(stderr, "Error: Failed to allocate root object\n");
        return TSL_PROGRAM_RESULT_ERR;
    }
    tsl_hash_map_init(self->variables);
    tsl_buffer_init(&self->stack_values);

    printf("###########################\n");
    result = tsl_program_run_function(self, 0);

    tsl_buffer_deinit(&self->stack_values);
    GC_FREE(self->variables); /* Free the root object, resulting in all objects in the program being free'd */
    self->variables = NULL;
    return result;
}