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

#include "ast.h"
#include "io.h"
#include "limits.h"
#include "module.h"
#include "parser.h"
#include "scanner.h"
#include "symtab.h"
#include "util.h"

/* Print an error string to `stderr`. */
#define error(...)                                                             \
    do {                                                                       \
        fprintf(stderr, "module: ");                                           \
        fprintf(stderr, __VA_ARGS__);                                          \
        fprintf(stderr, "\n");                                                 \
    } while (0)

/* Extract the directory part of a path */
static void getdirname(const char *path, char *result, size_t maxlen) {
    char buf[MAX_PATH_LEN];
    strndup(buf, path, MAX_PATH_LEN);

    char *dir = dirname(buf);
    strndup(result, dir, maxlen);
}

/* Check if a file has .rad extension */
static bool is_source(const char *path) {
    const char *dot = strrchr(path, '.');
    if (!dot)
        return false;
    return strcmp(dot, SOURCE_EXT) == 0;
}

/* Helper function to convert path with / to qualified name with :: */
static void path_to_qualified(const char *path, char *qualified, usize len) {
    usize j = 0;

    for (usize i = 0; path[i] && j < len - 2; i++) {
        if (path[i] == '/') {
            qualified[j++] = ':';
            qualified[j++] = ':';
        } else {
            qualified[j++] = path[i];
        }
    }
    qualified[j] = '\0';
}

/* Extract the file name without directory and extension */
static void getbasename(const char *path, char *result, size_t maxlen) {
    char buf[MAX_PATH_LEN];
    strndup(buf, path, MAX_PATH_LEN);

    /* Remove .rad extension if present */
    char *base = basename(buf);
    if (is_source(base)) {
        *strrchr(base, '.') = '\0';
    }
    strndup(result, base, maxlen);
}

/* Initialize the module manager */
void module_manager_init(module_manager_t *mm, const char *entryfile) {
    mm->nmodules = 0;
    getdirname(entryfile, mm->rootdir, MAX_PATH_LEN);
}

/* Initialize a module */
static void module_init(module_t *module, const char *path) {
    memset(module, 0, sizeof(*module));
    strndup(module->path, path, MAX_PATH_LEN);
    getbasename(path, module->name, MAX_PATH_LEN);
    strndup(module->qualified, module->name, MAX_PATH_LEN);

    module->state      = MODULE_STATE_UNVISITED;
    module->declared   = false;
    module->checked    = false;
    module->compiled   = false;
    module->source     = NULL;
    module->ast        = NULL;
    module->scope      = NULL;
    module->default_fn = NULL;
    module->parent     = NULL;
    module->nchildren  = 0;
}

/* Helper function to create a module path from a string */
void module_path(char dest[MAX_QUALIFIED_NAME], const char *name) {
    strncpy(dest, name, MAX_QUALIFIED_NAME);
}

/* Helper function to append a path component to a module path */
void module_qualify_str(
    char dest[MAX_QUALIFIED_NAME], const char *child, u16 len
) {
    strlcat(dest, "::", MAX_QUALIFIED_NAME);
    strncat(dest, child, len);
}

/* Helper function to append a path component to a module path */
void module_qualify(char dest[MAX_QUALIFIED_NAME], node_t *ident) {
    module_qualify_str(dest, ident->val.ident.name, ident->val.ident.length);
}

/* Add a module to the manager with custom qualified name */
module_t *module_manager_register_qualified(
    module_manager_t *mm, const char *path, const char *qualified
) {
    if (!mm || !path)
        return NULL;

    if (mm->nmodules >= MAX_MODULES) {
        error("maximum number of modules (%d) exceeded", MAX_MODULES);
        return NULL;
    }
    module_t *mod = NULL;
    if ((mod = module_manager_lookup(mm, path))) {
        return mod;
    }
    mod = &mm->modules[mm->nmodules++];
    module_init(mod, path);

    /* Build hierarchy if qualified name provided */
    if (qualified) {
        char path_copy[MAX_PATH_LEN];
        strncpy(path_copy, qualified, MAX_PATH_LEN - 1);
        path_copy[MAX_PATH_LEN - 1] = '\0';

        /* Remove .rad extension */
        char *ext = strrchr(path_copy, '.');
        if (ext)
            *ext = '\0';

        /* Check if this is a submodule (contains /) */
        char *slash = strrchr(path_copy, '/'); // Find LAST slash
        if (slash) {
            *slash            = '\0';
            char *parent_path = path_copy; // Everything before last slash
            char *child_name  = slash + 1; // Everything after last slash

            /* Convert parent path to qualified name for lookup */
            char parent_q[MAX_PATH_LEN];
            path_to_qualified(parent_path, parent_q, MAX_PATH_LEN);

            /* Look up parent using :: notation */
            module_t *parent =
                module_manager_lookup_by_qualified_name(mm, parent_q);

            if (!parent) {
                error(
                    "parent module '%s' not found for '%s'", parent_q, qualified
                );
                mm->nmodules--; // Rollback the module we just added
                return NULL;
            }
            mod->parent                           = parent;
            parent->children[parent->nchildren++] = mod;

            /* Write module qualified name */
            module_path(mod->qualified, parent_q);
            module_qualify_str(mod->qualified, child_name, strlen(child_name));
        } else {
            /* No parent (root-level module) */
            module_path(mod->qualified, path_copy);
        }
    }

    return mod;
}

/* Add a module to the manager */
module_t *module_manager_register(module_manager_t *mm, const char *path) {
    return module_manager_register_qualified(mm, path, NULL);
}

/* Find a module in the manager by path */
module_t *module_manager_lookup(module_manager_t *mm, const char *path) {
    for (usize i = 0; i < mm->nmodules; i++) {
        if (strcmp(mm->modules[i].path, path) == 0) {
            return &mm->modules[i];
        }
    }
    return NULL;
}

/* Find a module by name in the module manager */
module_t *module_manager_lookup_by_name(
    module_manager_t *mm, const char *name, u16 length
) {
    for (usize j = 0; j < mm->nmodules; j++) {
        if (strncmp(mm->modules[j].name, name, length) == 0 &&
            strlen(mm->modules[j].name) == length) {
            return &mm->modules[j];
        }
    }
    return NULL;
}

/* Find a module by qualified name in the module manager */
module_t *module_manager_lookup_by_qualified_name(
    module_manager_t *mm, const char *name
) {
    for (usize j = 0; j < mm->nmodules; j++) {
        if (strcmp(mm->modules[j].qualified, name) == 0) {
            return &mm->modules[j];
        }
    }
    return NULL;
}

/* Parse a single module, attaching diagnostics on failure. */
bool module_parse(module_t *module, i32 *err) {
    i32 size = readfile(module->path, &module->source);
    if (size < 0) {
        *err = MODULE_NOT_FOUND;
        error("unable to read module file at '%s'", module->path);
        return false;
    }
    scanner_init(&module->parser.scanner, module->path, module->source);
    parser_init(&module->parser);

    node_t *ast = parser_parse(&module->parser);
    if (module->parser.errors) {
        *err = MODULE_PARSE_ERROR;
        return false;
    }
    if (!ast) {
        *err = MODULE_PARSE_ERROR;
        error("failed to parse module '%s'", module->path);
        return false;
    }
    if (ast->cls != NODE_MOD_BODY) {
        *err = MODULE_PARSE_ERROR;
        error(
            "module '%s' parsed with unexpected root node (%d)",
            module->path,
            ast->cls
        );
        return false;
    }

    symbol_t sym = (symbol_t){
        .name   = module->name,
        .length = strlen(module->name),
        .node   = ast,
        .kind   = SYM_MODULE,
        .e.mod  = module,
    };
    module->ast      = ast;
    module->ast->sym = alloc_symbol(sym);

    return true;
}

bool module_manager_parse(module_manager_t *mm, i32 *err) {
    *err = MODULE_OK;

    for (usize i = 0; i < mm->nmodules; i++) {
        if (!(module_parse(&mm->modules[i], err))) {
            return false;
        }
    }
    return true;
}

/* Find module by relative import path */
module_t *module_manager_find_relative(
    module_manager_t *mm, const char *basepath, const char *import
) {
    /* Find the importing module */
    module_t *importer = module_manager_lookup(mm, basepath);
    if (!importer) {
        return NULL;
    }

    char import_copy[MAX_PATH_LEN];
    strndup(import_copy, import, MAX_PATH_LEN);

    char *segments[MAX_MODULES];
    usize nsegments = 0;

    for (char *token = strtok(import_copy, ":"); token;
         token       = strtok(NULL, ":")) {
        if (*token == '\0')
            continue;
        if (nsegments >= MAX_MODULES)
            break;
        segments[nsegments++] = token;
    }
    if (nsegments == 0)
        return NULL;

    module_t *current = importer;
    usize     index   = 0;

    while (index < nsegments && strcmp(segments[index], "super") == 0) {
        if (!current || !current->parent)
            return NULL;
        current = current->parent;
        index++;
    }

    if (index == nsegments)
        return current;

    if (index == 0 && nsegments == 1) {
        const char *segment = segments[0];

        for (usize i = 0; i < importer->nchildren; i++) {
            if (strcmp(importer->children[i]->name, segment) == 0) {
                return importer->children[i];
            }
        }
        return module_manager_lookup_by_name(mm, segment, strlen(segment));
    }

    if (index == 0) {
        current = module_manager_lookup_by_name(
            mm, segments[index], strlen(segments[index])
        );
        index++;
    }

    for (; index < nsegments && current; index++) {
        module_t *child = NULL;
        for (usize j = 0; j < current->nchildren; j++) {
            if (strcmp(current->children[j]->name, segments[index]) == 0) {
                child = current->children[j];
                break;
            }
        }
        current = child;
    }
    return current;
}