#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);
}

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, TslStackValue *src, TslValue *dst) {
    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: {
            TslStackValue *items_begin = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * src->data.integer);
            TslStackValue *items_end = items_begin + src->data.integer;
            TslStackValue *item = items_end - 1;
            TslValue *list_item;
            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(item != items_begin - 1) {
                return_if_error(tsl_value_create_from_stack_value(self, item, list_item));
                --list_item;
                --item;
            }

            dst->type = TSL_TYPE_LIST;
            break;
        }
        case TSL_STACK_VALUE_TYPE_MAP: {
            TslStackValue *items_begin = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * src->data.integer);
            TslStackValue *items_end = items_begin + src->data.integer;
            TslStackValue *item = items_end - 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(item != items_begin - 1) {
                TslValue map_key;
                TslValue map_value;
                return_if_error(tsl_value_create_from_stack_value(self, item, &map_value));
                return_if_error(tsl_value_create_from_stack_value(self, item - 1, &map_key));
                return_if_error(tsl_hash_map_insert(dst->data.map, &map_key, &map_value));
                item -= 2;
            }

            dst->type = TSL_TYPE_MAP;
            break;
        }
        case TSL_STACK_VALUE_TYPE_INDEX: {
            TslStackValue *args = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * 2);
            TslValue var;
            TslValue key;

            return_if_error(tsl_value_create_from_stack_value(self, args + 1, &key));
            return_if_error(tsl_value_create_from_stack_value(self, args, &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;
}

TslProgramResult tsl_program_run(TslProgram *self) {
    TslProgramResult result = TSL_PROGRAM_RESULT_OK;
    TslBytecode *file_scope_bytecode = tsl_buffer_begin(&self->function_bytecode_list);
    char *instruction = tsl_buffer_begin(&file_scope_bytecode->buffer);
    char *instruction_end = tsl_buffer_end(&file_scope_bytecode->buffer);

    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");

    /* 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: {
                TslStackValue *stack_value = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue));
                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, stack_value, 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: {
                /* + 1, since the first arg is the function index and the other args are the args for the function */
                TslStackValue *args = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * (1 + instruction_type1->value));
                (void)args;
                /* TODO: Implement this */
                printf("callf %d\n", instruction_type1->value);
                instruction += sizeof(TslInstructionType1);
                break;
            }
            case TSL_OPCODE_ADD: {
                TslStackValue *args = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * 2);
                TslValue lhs_value;
                TslValue rhs_value;
                TslStackValue stack_value;

                cleanup_if_error(tsl_value_create_from_stack_value(self, args + 1, &rhs_value));
                cleanup_if_error(tsl_value_create_from_stack_value(self, args, &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: {
                TslStackValue *args = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * 2);
                (void)args;
                /* TODO: Implement this */
                printf("sub\n");
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_MUL: {
                TslStackValue *args = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * 2);
                (void)args;
                /* TODO: Implement this */
                printf("mul\n");
                instruction += sizeof(TslInstructionType5);
                break;
            }
            case TSL_OPCODE_DIV: {
                TslStackValue *args = tsl_buffer_pop(&self->stack_values, sizeof(TslStackValue) * 2);
                (void)args;
                /* TODO: Implement this */
                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;
                }
                
                printf("callc %d\n", instruction_type1->value);
                instruction += sizeof(TslInstructionType1);
                break;
            }
        }
    }

    cleanup:
    assert(self->stack_values.size == 0); /* All push instructions (mostly load) haven't been handled yet. This is a bug in the compiler */
    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;
}