#include "../include/compiler.h"
#include "../include/parser.h"
#include "../include/std/log.h"
#include "../include/std/mem.h"
#include "../include/std/alloc.h"
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>

#define MIN(a, b) ((a) < (b) ? (a) : (b))

static CHECK_RESULT int get_thread_count_env_var(int *thread_count) {
    char *threads;
    threads = getenv("THREADS");
    if(!threads)
        return -1;
    *thread_count = atoi(threads);
    return 0;
}

int amal_compiler_init(amal_compiler *self) {
    int i;
    int result;

    result = get_thread_count_env_var(&self->usable_thread_count);
    if(result != 0) {
        self->usable_thread_count = amal_get_usable_thread_count();
        if(self->usable_thread_count == 0) {
            amal_log_info("Unable to get the number of threads available on the system, using 1 thread.");
            amal_log_info("You can override the number of threads using by setting the environment variable THREADS");
            self->usable_thread_count = 1;
        }
    } else if(self->usable_thread_count == 0) {
        amal_log_info("Environment variable THREADS contains invalid number for threads. THREADS has to be at least 1.");
        return AMAL_COMPILER_ERR;
    }

    am_memset(&self->allocator, 0, sizeof(self->allocator));
    am_memset(&self->main_thread_allocator, 0, sizeof(self->main_thread_allocator));
    self->started = bool_false;
    amal_mutex_init(&self->mutex);

    return_if_error(scoped_allocator_init(&self->allocator));
    cleanup_if_error(scoped_allocator_init(&self->main_thread_allocator));
    cleanup_if_error(buffer_init(&self->parsers, &self->allocator));
    cleanup_if_error(buffer_init(&self->queued_files, &self->allocator));
    cleanup_if_error(scoped_allocator_alloc(&self->allocator,
                                            self->usable_thread_count * sizeof(ParserThreadData),
                                            (void**)&self->threads));
    for(i = 0; i < self->usable_thread_count; ++i)
        cleanup_if_error(parser_thread_data_init(&self->threads[i]));
    return AMAL_COMPILER_OK;

    cleanup:
    /* Ignore result */
    result = amal_compiler_deinit(self);
    return AMAL_COMPILER_ERR;
}

int amal_compiler_deinit(amal_compiler *self) {
    int i;
    int result;
    result = AMAL_COMPILER_OK;

    for(i = 0; i < self->usable_thread_count; ++i) {
        int r;
        r = parser_thread_data_deinit(&self->threads[i]);
        if(r != 0)
            result = r;
    }

    scoped_allocator_deinit(&self->allocator);
    scoped_allocator_deinit(&self->main_thread_allocator);
    return result;
}

typedef struct {
    amal_compiler *compiler;
    ParserThreadData *parser_thread_data;
    BufferView filepath;
} CompilerParserThreadUserData;

static CHECK_RESULT int amal_compiler_load_first_this_thread(amal_compiler *self, BufferView filepath, ScopedAllocator *allocator) {
    Parser parser;
    int result;
    result = AMAL_COMPILER_ERR;
    am_memset(&parser, 0, sizeof(parser));

    return_if_error(parser_init(&parser, self, allocator));
    cleanup_if_error(parser_parse_file(&parser, filepath));
    amal_log_debug("Finished parsing in thread");
    cleanup_if_error(amal_mutex_lock(&self->mutex, "amal_compiler_load_first_this_thread"));
    cleanup_if_error(buffer_append(&self->parsers, &parser, sizeof(parser)));
    result = AMAL_COMPILER_OK;

    cleanup:
    amal_mutex_tryunlock(&self->mutex);
    return result;
}

/* TODO: Handle errors (stop parsing in all other threads and report errors/warnings) */
static void* thread_callback_parse_file(void *userdata) {
    int _;
    BufferView next_file;
    CompilerParserThreadUserData compiler_parser_userdata;
    (void)_;
    next_file = create_buffer_view_null();
    assert(!amal_thread_is_main());
    
    am_memcpy(&compiler_parser_userdata, userdata, sizeof(compiler_parser_userdata));
    am_free(userdata);
    amal_log_debug("Thread start parse file");
    cleanup_if_error(amal_compiler_load_first_this_thread(compiler_parser_userdata.compiler, 
                                                          compiler_parser_userdata.filepath,
                                                          &compiler_parser_userdata.parser_thread_data->allocator));
    cleanup_if_error(amal_mutex_lock(&compiler_parser_userdata.compiler->mutex, "thread_callback_parse_file"));
    if(compiler_parser_userdata.compiler->queued_files.size > 0)
        _ = buffer_pop(&compiler_parser_userdata.compiler->queued_files, &next_file, sizeof(next_file));

    cleanup:
    amal_log_debug("Thread finished, next file: %s", next_file.data ? "yes" : "no");
    if(!next_file.data)
        compiler_parser_userdata.parser_thread_data->status = PARSER_THREAD_STATUS_IDLE;
    amal_mutex_tryunlock(&compiler_parser_userdata.compiler->mutex);
    if(next_file.data) {
        _ = amal_compiler_load_first_this_thread(compiler_parser_userdata.compiler,
                                                 next_file,
                                                 &compiler_parser_userdata.parser_thread_data->allocator);
    }
    return NULL;
}

static CHECK_RESULT int amal_compiler_load_file_select_thread(amal_compiler *self, BufferView filepath, ParserThreadData **thread_selected) {
    int i;
    int result;
    ParserThreadData *parser_thread_data;
    CompilerParserThreadUserData *thread_user_data;
    thread_user_data = NULL;
    *thread_selected = NULL;
    result = AMAL_COMPILER_ERR;

    cleanup_if_error(amal_mutex_lock(&self->mutex, "amal_compiler_load_file_select_thread"));
    for(i = 0; i < self->usable_thread_count; ++i) {
        parser_thread_data = &self->threads[i];
        if(parser_thread_data->status == PARSER_THREAD_STATUS_RUNNING)
            continue;

        cleanup_if_error(am_malloc(sizeof(CompilerParserThreadUserData), (void**)&thread_user_data));
        thread_user_data->compiler = self;
        thread_user_data->parser_thread_data = parser_thread_data;
        thread_user_data->filepath = filepath;
        result = parser_thread_data_start(parser_thread_data, thread_callback_parse_file, thread_user_data);
        *thread_selected = parser_thread_data;
        break;
    }

    cleanup:
    if(result != 0)
        am_free(thread_user_data);
    amal_mutex_tryunlock(&self->mutex);
    return result;
}

static CHECK_RESULT int amal_compiler_load_file_join_threads(amal_compiler *self) {
    int i;
    int _;
    int result;
    void *thread_return_data;
    ParserThreadData *parser_thread_data;
    result = AMAL_COMPILER_ERR;
    (void)_;
    assert(amal_thread_is_main());
    
    for(i = 0; i < self->usable_thread_count; ++i) {
        cleanup_if_error(amal_mutex_lock(&self->mutex, "amal_compiler_load_file_join_threads, waiting for workers"));
        parser_thread_data = &self->threads[i];
        amal_mutex_tryunlock(&self->mutex);
        amal_log_debug("Joining thread %d, status %d", i, parser_thread_data->status);
        _ = parser_thread_data_join(parser_thread_data, &thread_return_data);
    }

    result = AMAL_COMPILER_OK;
    cleanup:
    return result;
}

int amal_compiler_load_file(amal_compiler *self, BufferView filepath) {
    int result;
    ParserThreadData *parser_thread_data;
    result = AMAL_COMPILER_ERR;

    return_if_error(amal_compiler_load_file_select_thread(self, filepath, &parser_thread_data));

    /*
    amal_compiler_load_file is called by the user for the first file to compile
    but also called by the parser when it sees @import
    */
    if(!self->started) {
        self->started = bool_true;
        /*amal_log_info("Parsing %.*s using %d thread(s)", (int)filepath.size, filepath.data, self->usable_thread_count);*/
        /*return_if_error(amal_compiler_load_first_this_thread(self, filepath, &self->main_thread_allocator));*/
        return amal_compiler_load_file_join_threads(self);
    }

    if(parser_thread_data)
        return AMAL_COMPILER_OK;

    cleanup_if_error(amal_mutex_lock(&self->mutex, "amal_compiler_load_file"));
    cleanup_if_error(buffer_append(&self->queued_files, &filepath, sizeof(filepath)));
    result = AMAL_COMPILER_OK;

    cleanup:
    amal_mutex_tryunlock(&self->mutex);
    return result;
}