//
// Created by servostar on 6/2/24.
//

#include <assert.h>
#include <ast/ast.h>
#include <cfg/opt.h>
#include <codegen/backend.h>
#include <compiler.h>
#include <glib.h>
#include <io/files.h>
#include <lex/util.h>
#include <link/lib.h>
#include <llvm/backend.h>
#include <mem/cache.h>
#include <set/set.h>
#include <stdlib.h>
#include <sys/log.h>
#include <unistd.h>
#include <yacc/parser.tab.h>

#define GRAPHVIZ_FILE_EXTENSION "gv"

extern void yyrestart(FILE*);

// Module AST node used by the parser for AST construction.
[[maybe_unused]]
AST_NODE_PTR root;
// Current file which gets compiled the parser.
// NOTE: due to global state no concurrent compilation is possible
//       on parser level.
[[maybe_unused]]
ModuleFile* current_file;

[[maybe_unused]]
ModuleRef* parser_ref;
ModuleRef* module_ref;

static int build_project_targets(ModuleFileStack* unit,
                                 const ProjectConfig* config);

static int build_target(ModuleFileStack* unit, const TargetConfig* target);

/**
 * @brief Compile the specified file into AST
 * @param ast Initialized AST module node to build program rules
 * @param file The file to be processed
 * @return EXIT_SUCCESS in case the parsing was success full anything lese if
 * not
 */
[[nodiscard("AST may be in invalid state")]] [[gnu::nonnull(1),
                                               gnu::nonnull(1)]]
static int compile_file_to_ast(AST_NODE_PTR ast, ModuleFile* file) {
    assert(file->path != NULL);
    assert(ast != NULL);

    file->handle = fopen(file->path, "r");

    if (file->handle == NULL) {
        INFO("unable to open file: %s", file->path);
        print_message(Error, "Cannot open file %s: %s", file->path,
                      strerror(errno));
        return EXIT_FAILURE;
    }

    DEBUG("parsing file: %s", file->path);
    // setup global state
    root         = ast;
    current_file = file;
    yyin         = file->handle;
    parser_ref   = module_ref_clone(module_ref);
    yyrestart(yyin);
    lex_reset();

    int status = yyparse();

    // clean up global state
    // current_file = NULL;
    root = NULL;
    yyin = NULL;

    return status;
}

/**
 * @brief Setup the environment of the target.
 * @param target
 * @return EXIT_SUCCESS if successful EXIT_FAILURE otherwise
 */
static int setup_target_environment(const TargetConfig* target) {
    DEBUG("setting up environment for target: %s", target->name);

    assert(target->output_directory != NULL);
    assert(target->archive_directory != NULL);

    int result = create_directory(target->archive_directory);
    if (result != 0 && errno != EEXIST) {
        const char* message = get_last_error();
        assert(message != NULL);

        print_message(Error, "Unable to create directory: %s: %s",
                      target->archive_directory, message);
        free((void*) message);
        return EXIT_FAILURE;
    }

    result = create_directory(target->output_directory);
    if (result != 0 && errno != EEXIST) {
        const char* message = get_last_error();
        assert(message != NULL);

        print_message(Error, "Unable to create directory: %s: %s",
                      target->output_directory, message);
        free((void*) message);
        return EXIT_FAILURE;
    }

    INFO("setup environment successfully");

    return EXIT_SUCCESS;
}

/**
 * @brief Print the supplied AST of the specified target to a graphviz ".gv"
 * file
 * @param ast
 * @param target
 */
static void print_ast_to_file(AST_NODE_PTR ast, const TargetConfig* target) {
    assert(ast != NULL);
    assert(target != NULL);
    DEBUG("printing AST to file: %s", target->name);

    if (!target->print_ast) {
        INFO("no need to print AST");
        return;
    }

    // create file path to write graphviz to
    // basename of ile
    char* filename =
      g_strjoin(".", target->name, GRAPHVIZ_FILE_EXTENSION, NULL);
    // relative path to file
    char* path = g_build_filename(target->archive_directory, filename, NULL);

    DEBUG("Opening file to graph: %s", path);

    FILE* output = fopen(path, "w");
    if (output == NULL) {
        char* message = (char*) get_last_error();
        print_message(Error, "Unable to open file for syntax tree at: %s: %s",
                      path, message);
        free(message);
    } else {
        DEBUG("writing graph to file...");

        AST_fprint_graphviz(output, ast);

        fclose(output);

        print_message(Info, "AST graph was written to: %s", path);
    }

    g_free(filename);
    g_free(path);
}

static int run_backend_codegen(const Module* module,
                               const TargetConfig* target) {
    DEBUG("initializing LLVM codegen backend...");
    llvm_backend_init();

    DEBUG("initiializing backend for codegen...");
    BackendError err = init_backend();
    if (err.kind != Success) {
        return EXIT_FAILURE;
    }

    DEBUG("generating code...");
    err = generate_code(module, target);
    if (err.kind != Success) {
        print_message(Error, "Backend failed: %s", err.impl.message);
        return EXIT_FAILURE;
    }

    print_message(Info, "Compilation finished successfully");

    err = deinit_backend();
    if (err.kind != Success) {
        ERROR("Unable to deinit backend: %s", err.impl.message);
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}

const char* get_absolute_import_path(const TargetConfig* config,
                                     const char* import_target_name) {
    INFO("resolving absolute path for import target: %s", import_target_name);

    if (!g_str_has_suffix(import_target_name, ".gsc")) {
        char* full_filename = g_strjoin("", import_target_name, ".gsc", NULL);
        import_target_name  = mem_strdup(MemoryNamespaceLld, full_filename);
        g_free(full_filename);
    }

    for (guint i = 0; i < config->import_paths->len; i++) {
        const char* import_directory_path =
          g_array_index(config->import_paths, char*, i);

        char* path =
          g_build_filename(import_directory_path, import_target_name, NULL);
        char* cwd       = g_get_current_dir();
        char* canonical = g_canonicalize_filename(path, cwd);

        const gboolean exists = g_file_test(canonical, G_FILE_TEST_EXISTS);
        const gboolean is_dir = g_file_test(canonical, G_FILE_TEST_IS_DIR);

        char* cached_canonical = mem_strdup(MemoryNamespaceLld, canonical);

        g_free(path);
        g_free(cwd);
        g_free(canonical);

        if (exists && !is_dir) {
            INFO("import target found at: %s", cached_canonical);
            return cached_canonical;
        }
    }

    // file not found
    return NULL;
}

static int compile_module_with_dependencies(ModuleFileStack* unit,
                                            ModuleFile* file,
                                            const TargetConfig* target,
                                            AST_NODE_PTR root_module) {

    if (NULL == module_ref) {
        module_ref = mem_alloc(MemoryNamespaceOpt, sizeof(ModuleRef));
        module_ref->module_path =
          mem_new_g_array(MemoryNamespaceOpt, sizeof(char*));
    } else {
        module_ref_push(module_ref, module_from_basename(file->path));
    }

    GHashTable* imports =
      mem_new_g_hash_table(MemoryNamespaceAst, g_str_hash, g_str_equal);

    if (compile_file_to_ast(root_module, file) == EXIT_SUCCESS) {

        for (size_t i = 0; i < AST_get_child_count(root_module); i++) {
            AST_NODE_PTR child = AST_get_node(root_module, i);

            if (child->kind == AST_Import) {
                if (g_hash_table_contains(target->dependencies, child->value)) {
                    Dependency* dependency =
                      g_hash_table_lookup(target->dependencies, child->value);

                    switch (dependency->kind) {
                        case GemstoneProject:
                            gchar* cwd = g_get_current_dir();
                            chdir(dependency->mode.project.path);

                            ProjectConfig* new_config =
                              default_project_config();
                            if (load_project_config(new_config)) {
                                print_message(
                                  Error, "Failed to load project config: `%s`",
                                  child->value);
                                return EXIT_FAILURE;
                            }

                            TargetConfig* dep_target = g_hash_table_lookup(
                              new_config->targets,
                              dependency->mode.project.target);
                            if (build_target(unit, dep_target)) {
                                print_message(
                                  Error, "Failed to build project config: `%s`",
                                  child->value);
                                return EXIT_FAILURE;
                            }

                            GPathBuf* buf = g_path_buf_new_from_path(
                              dependency->mode.project.path);
                            TargetConfig* dep_conf = g_hash_table_lookup(
                              new_config->targets,
                              dependency->mode.project.target);
                            char* root_mod = dep_conf->root_module;
                            g_path_buf_push(buf, root_mod);
                            char* rel_path = g_path_buf_to_path(buf);

                            GPathBuf* dep_bin = g_path_buf_new();
                            g_path_buf_push(dep_bin,
                                            dependency->mode.project.path);
                            g_path_buf_push(dep_bin,
                                            dep_conf->archive_directory);
                            char* dep_obj_file =
                              g_strjoin(".", dep_conf->name, "o", NULL);
                            g_path_buf_push(dep_bin, dep_obj_file);
                            char* dep_bin_path =
                              g_path_buf_free_to_path(dep_bin);
                            g_free(dep_obj_file);
                            char* cached_dep_bin_path =
                              mem_strdup(MemoryNamespaceLld, dep_bin_path);
                            g_free(dep_bin_path);

                            g_array_append_val(dependency->libraries,
                                               cached_dep_bin_path);

                            const char* path =
                              get_absolute_import_path(target, rel_path);
                            g_free(rel_path);

                            if (g_hash_table_contains(imports, path)) {
                                continue;
                            }

                            ModuleFile* imported_file = push_file(unit, path);
                            AST_NODE_PTR imported_module = AST_new_node(
                              empty_location(imported_file, module_ref),
                              AST_Module, NULL);

                            if (compile_module_with_dependencies(
                                  unit, imported_file, dep_conf,
                                  imported_module)
                                == EXIT_SUCCESS) {
                                AST_import_module(root_module, i + 1,
                                                  imported_module);
                            } else {
                                return EXIT_FAILURE;
                            }

                            g_hash_table_insert(imports, (gpointer) path, NULL);

                            g_path_buf_pop(buf);
                            gchar* directory = g_path_buf_free_to_path(buf);
                            gchar* cached_directory =
                              mem_strdup(MemoryNamespaceLld, directory);
                            g_free(directory);
                            g_array_append_val(target->import_paths,
                                               cached_directory);

                            chdir(cwd);
                            g_free(cwd);

                            GHashTableIter iter;

                            g_hash_table_iter_init(&iter,
                                                   dep_target->dependencies);
                            char* key;
                            Dependency* dep;
                            while (g_hash_table_iter_next(
                              &iter, (gpointer) &key, (gpointer) &dep)) {
                                if (dep->kind == GemstoneProject) {
                                    for (guint i = 0; i < dep->libraries->len;
                                         i++) {
                                        char* dep_lib = g_array_index(
                                          dep->libraries, char*, i);
                                        g_array_append_val(
                                          dependency->libraries, dep_lib);
                                    }
                                } else if (dep->kind == NativeLibrary) {
                                    char* library_name =
                                      build_platform_library_name(
                                        dep->mode.library.name,
                                        dep->mode.library.shared);
                                    g_array_append_val(dependency->libraries,
                                                       library_name);
                                }
                            }

                            break;
                        case NativeLibrary:

                            char* library_name = build_platform_library_name(
                              dependency->mode.library.name,
                              dependency->mode.library.shared);
                            g_array_append_val(dependency->libraries,
                                               library_name);

                            break;
                        default:
                            break;
                    }

                } else {
                    print_message(Error, "Cannot resolve path for import: `%s`",
                                  child->value);
                    return EXIT_FAILURE;
                }

            } else if (child->kind == AST_Include) {

                const char* path =
                  get_absolute_import_path(target, child->value);
                if (path == NULL) {
                    print_message(Error,
                                  "Cannot resolve path for include: `%s`",
                                  child->value);
                    return EXIT_FAILURE;
                }

                if (g_hash_table_contains(imports, path)) {
                    continue;
                }

                module_ref_push(module_ref, module_from_basename(path));

                ModuleFile* imported_file    = push_file(unit, path);
                AST_NODE_PTR imported_module = AST_new_node(
                  empty_location(imported_file, module_ref), AST_Module, NULL);

                if (compile_file_to_ast(imported_module, imported_file)
                    == EXIT_SUCCESS) {
                    AST_merge_modules(root_module, i + 1, imported_module);
                } else {
                    return EXIT_FAILURE;
                }

                g_hash_table_insert(imports, (gpointer) path, NULL);

                gchar* directory = g_path_get_dirname(path);
                gchar* cached_directory =
                  mem_strdup(MemoryNamespaceLld, directory);
                g_free(directory);
                g_array_append_val(target->import_paths, cached_directory);

                module_ref_pop(module_ref);
            }
        }
    } else {
        return EXIT_FAILURE;
    }

    module_ref_pop(module_ref);

    return EXIT_SUCCESS;
}

/**
 * @brief Build the given target
 * @param unit
 * @param target
 */
static int build_target(ModuleFileStack* unit, const TargetConfig* target) {
    int err = EXIT_SUCCESS;

    print_message(Info, "Building target: %s", target->name);

    ModuleFile* file = push_file(unit, target->root_module);
    AST_NODE_PTR root_module =
      AST_new_node(empty_location(file, module_ref), AST_Module, NULL);

    err = compile_module_with_dependencies(unit, file, target, root_module);
    if (err == EXIT_SUCCESS) {
        if (root_module != NULL) {
            err = setup_target_environment(target);
            if (err == 0) {

                print_ast_to_file(root_module, target);
                Module* module = create_set(root_module);

                if (module != NULL) {
                    err = run_backend_codegen(module, target);
                } else {
                    err = EXIT_FAILURE;
                }
            }

            AST_delete_node(root_module);
        }
    }

    print_file_statistics(file);

    return err;
}

/**
 * @brief Compile a single file.
 *        Creates a single target by the given command line arguments.
 * @param unit
 */
static int compile_file(ModuleFileStack* unit) {
    INFO("compiling basic files...");

    TargetConfig* target = default_target_config_from_args();

    if (target->root_module == NULL) {
        print_message(Error, "No input file specified.");
        delete_target_config(target);
        return EXIT_FAILURE;
    }

    int err = build_target(unit, target);

    delete_target_config(target);

    return err;
}

/**
 * @brief Build all project targets specified by the command line arguments.
 * @param unit
 * @param config
 */
static int build_project_targets(ModuleFileStack* unit,
                                 const ProjectConfig* config) {
    int err = EXIT_SUCCESS;

    if (is_option_set("all")) {
        // build all targets in the project
        GHashTableIter iter;

        g_hash_table_iter_init(&iter, config->targets);

        char* key;
        TargetConfig* val;
        while (
          g_hash_table_iter_next(&iter, (gpointer) &key, (gpointer) &val)) {
            err = build_target(unit, val);
        }
        return err;
    }

    // build all targets given in the arguments
    GArray* targets = get_non_options_after("build");

    if (targets != NULL) {
        for (guint i = 0; i < targets->len; i++) {
            const char* target_name = g_array_index(targets, const char*, i);

            if (g_hash_table_contains(config->targets, target_name)) {
                err = build_target(
                  unit, g_hash_table_lookup(config->targets, target_name));
            } else {
                print_message(Error, "Unknown target: %s", target_name);
            }
        }

        mem_free(targets);
    } else {
        print_message(Error, "No targets specified.");
    }

    return err;
}

/**
 * @brief Build targets from project. Configuration is provided by command line
 * arguments.
 * @param unit File storage
 */
static int build_project(ModuleFileStack* unit) {
    ProjectConfig* config = default_project_config();
    int err               = load_project_config(config);

    if (err == PROJECT_OK) {
        err = build_project_targets(unit, config);
    }

    delete_project_config(config);
    return err;
}

int run_compiler() {
    ModuleFileStack files = new_file_stack();

    int status = EXIT_SUCCESS;

    if (is_option_set("build")) {
        status = build_project(&files);
    } else if (is_option_set("compile")) {
        status = compile_file(&files);
    } else {
        print_message(Error, "Invalid mode of operation. Rerun with --help.");
    }

    if (files.files == NULL) {
        print_message(Error, "No input files, nothing to do.");
    } else {
        print_unit_statistics(&files);
    }

    delete_files(&files);

    return status;
}