/* splc0.c - SPL compiler CLI (stage 1, 引导用)
 *
 *   splc0 --dump tokens|ast|sema|ir|vm <file>    dump 前端产物
 *   splc0 <in> <out>                             编译
 */
#define __SCC_LOG_IMPL_IMPORT_SRC__
#include "../stage0/include/utils.h"

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

#include "spl_ast.h"
#include "spl_ast2ir.h"
#include "spl_diag.h"
#include "spl_ir2vm.h"
#include "spl_lexer.h"
#include "spl_parser.h"
#include "spl_sema.h"
#include "spl_tok.h"

typedef enum {
    DUMP_NONE,
    DUMP_TOKEN,
    DUMP_AST,
    DUMP_SEMA,
    DUMP_IR,
    DUMP_VM,
} dump_t;

typedef struct {
    spl_ast_t *ast;
    char *abs_path;
    char *src;
} spl_driver_module_t;
typedef VEC(spl_driver_module_t) spl_driver_module_vec_t;

typedef struct {
    spl_ast_t main_ast;
    spl_tok_vec_t main_toks;
    char *main_path;

    spl_driver_module_vec_t modules;
    MAP(const char *, usize) path_to_idx;

    spl_diag_t *diag;
} spl_driver_t;

static char *read_file(const char *path, long *out_len) {
    FILE *f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "path `%s` ", path);
        perror("fopen");
        return NULL;
    }
    fseek(f, 0, SEEK_END);
    long len = ftell(f);
    fseek(f, 0, SEEK_SET);
    char *buf = malloc((size_t)len + 1);
    if (!buf) {
        fclose(f);
        return NULL;
    }
    size_t nread = fread(buf, 1, (size_t)len, f);
    fclose(f);
    buf[nread] = '\0';
    if (out_len)
        *out_len = (long)nread;
    return buf;
}

/* 注意：按值接收，内部 vec_free 释放与调用方共享的 data 指针。
   调用后原 vec 视为已消费，不要再使用。 */
static void dump_tokens(spl_tok_vec_t toks) {
    printf("tokens got (%zu)\n", toks.size);
    for (usize i = 0; i < toks.size; i++) {
        const spl_tok_t *t = &toks.data[i];
        printf("[%s] %.*s (%zu:%zu)\n", spl_tok_type_name(t->type), (int)t->len, t->lexeme, t->line,
               t->col);
    }
    vec_free(toks);
}

/* 把 import_path 拼接到 importing_file 所在目录。
   from_fname 为空时按相对 cwd 处理。绝对路径直接透传。 */
static void resolve_relative(const char *from, const char *path, char *out, usize cap) {
    if (!from || !path) {
        if (cap)
            out[0] = 0;
        return;
    }
    if (path[0] == '/' || path[0] == '\\' || (path[0] && path[1] == ':')) {
        snprintf(out, cap, "%s", path);
        return;
    }
    const char *last_sep = NULL;
    for (const char *p = from; *p; ++p)
        if (*p == '/' || *p == '\\')
            last_sep = p;
    if (!last_sep) {
        snprintf(out, cap, "%s", path);
        return;
    }
    usize dir_len = (usize)(last_sep - from) + 1;
    snprintf(out, cap, "%.*s%s", (int)dir_len, from, path);
}

static spl_ast_t *driver_load(void *user, const char *from_fname, const char *path,
                              spl_diag_t *diag) {
    spl_driver_t *d = user;
    char abs[1024];
    resolve_relative(from_fname, path, abs, sizeof abs);

    usize idx;
    if (map_get(d->path_to_idx, abs, &idx)) {
        spl_ast_t *a = vec_at(d->modules, idx).ast;
        if (!a) {
            spl_diag_error(diag, SPL_DIAG_KIND_CIRCULAR_DEP, NULL, "circular import: `%s`", abs);
            return NULL;
        }
        return a;
    }

    /* 先登记 ast=NULL 占位：若后续解析过程中又有 @import 指回本路径，
       会命中这里，判定为环。 */
    char *key = strdup(abs);
    idx = vec_size(d->modules);
    spl_driver_module_t placeholder = {.ast = NULL, .abs_path = key, .src = NULL};
    vec_push(d->modules, placeholder);
    map_put(d->path_to_idx, key, idx);

    long len = 0;
    char *src = read_file(abs, &len);
    if (!src) {
        spl_diag_error(diag, SPL_DIAG_KIND_LEX_UNEXPECTED, NULL, "@import: cannot open `%s`", abs);
        return NULL;
    }
    (void)len;

    /* key 同时用作 lexer 的 fname：token.fname 与 map key 共享一份 strdup，
       减少一次分配；cleanup 时只 free module.abs_path */
    spl_tok_vec_t toks = spl_lex(src, key);
    spl_ast_t *ast = malloc(sizeof *ast);
    spl_ast_init(ast, &toks);
    spl_parser_t parser;
    spl_parser_init(&parser);
    parser.diag = diag;
    spl_ast_parse(&parser, ast);
    spl_parser_drop(&parser);

    vec_at(d->modules, idx).ast = ast;
    vec_at(d->modules, idx).src = src;
    return ast;
}

static int compile_spl(const char *src, const char *fname, const char *outpath, int gen_debug,
                       dump_t dump) {
    (void)gen_debug;

    spl_diag_t diag;
    spl_diag_init(&diag, stderr);

    spl_driver_t drv = {0};
    drv.diag = &diag;
    drv.main_path = strdup(fname);
    drv.main_toks = spl_lex(src, drv.main_path);
    spl_ast_init(&drv.main_ast, &drv.main_toks);
    map_init(drv.path_to_idx, MAP_HASH_STR, MAP_CMP_STR);
    vec_init(drv.modules);

    int rc = 0;

    /* dump_tokens 消费 drv.main_toks（共享 data），之后 main_ast.input 悬空，
       所以这条分支必须在 spl_ast_parse 之前、且不能再 drop main_ast。 */
    if (dump == DUMP_TOKEN) {
        dump_tokens(drv.main_toks);
        free(drv.main_path);
        map_free(drv.path_to_idx);
        spl_diag_drop(&diag);
        return 0;
    }

    spl_attr_t attr;
    spl_lang_type_t lang_type;
    spl_layout_target_t layout_target;
    spl_layout_t layout;
    spl_scope_t scope;
    spl_sema_t sema;
    spl_ast2ir_t a2ir;
    spl_ir_builder_t ir_builder;
    spl_ir2vm_t ir2vm;
    spl_ast_ptr_vec_t all_asts;

    int attr_ok = 0, lang_type_ok = 0, layout_ok = 0, scope_ok = 0;
    int sema_ok = 0, a2ir_ok = 0, irb_ok = 0, ir2vm_ok = 0;

    spl_diag_set_stage(&diag, SPL_STAGE_PARSE);
    spl_parser_t parser;
    spl_parser_init(&parser);
    parser.diag = &diag;
    spl_ast_parse(&parser, &drv.main_ast);
    spl_parser_drop(&parser);
    if (spl_diag_has_error(&diag) || spl_diag_has_fatal(&diag)) {
        spl_diag_summary(&diag);
        rc = 1;
        goto cleanup;
    }
    if (dump == DUMP_AST) {
        spl_ast_dump(&drv.main_ast, drv.main_ast.root);
        goto cleanup;
    }

    spl_attr_init(&attr);
    attr_ok = 1;
    spl_lang_type_init(&lang_type);
    lang_type_ok = 1;

    /* 目标描述符：剥离目标机器细节，layout 是 size/align/offset 的唯一来源。 */
    memset(&layout_target, 0, sizeof layout_target);
    layout_target.endian = SPL_LITTLE_ENDIAN;
    layout_target.ptr_size = sizeof(void *);
    layout_target.ptr_align = sizeof(void *);
    spl_layout_init(&layout, &layout_target);
    layout_ok = 1;

    spl_scope_init(&scope);
    scope_ok = 1;

    spl_sema_init(&sema, &drv.main_ast, &lang_type, &layout, &scope, &attr);
    sema_ok = 1;
    spl_module_loader_t loader = {.load = driver_load, .user = &drv};
    sema.loader = &loader;
    spl_diag_set_stage(&diag, SPL_STAGE_SEMA);
    sema.diag = &diag;
    spl_sema_run(&sema);
    if (spl_diag_has_error(&diag) || spl_diag_has_fatal(&diag)) {
        spl_diag_summary(&diag);
        rc = 1;
        goto cleanup;
    }
    if (dump == DUMP_SEMA) {
        spl_sema_dump(&sema);
        spl_ast_dump(&drv.main_ast, drv.main_ast.root);
        goto cleanup;
    }

    /* 收集所有 ast：主文件在前，被 import 的文件按加载顺序在后。
       ast2ir 内部会 vec_free(all_asts)。 */
    vec_init(all_asts);
    vec_push(all_asts, &drv.main_ast);
    for (usize i = 0; i < vec_size(drv.modules); ++i)
        vec_push(all_asts, vec_at(drv.modules, i).ast);

    spl_ir_builder_init(&ir_builder);
    irb_ok = 1;
    spl_ast2ir_init(&a2ir, &ir_builder, &layout, &sema);
    a2ir_ok = 1;
    spl_diag_set_stage(&diag, SPL_STAGE_AST2IR);
    a2ir.diag = &diag;
    spl_ast2ir_run(&a2ir, all_asts);
    if (spl_diag_has_error(&diag) || spl_diag_has_fatal(&diag)) {
        spl_diag_summary(&diag);
        rc = 1;
        goto cleanup;
    }
    if (dump == DUMP_IR) {
        spl_ir_dump(&ir_builder.ir, &lang_type.ir, &attr);
        goto cleanup;
    }

    /* 类型已全部就绪：解析 layout，供 ir2vm 的 size/align/offset 使用。 */
    spl_layout_resolve(&layout, &lang_type.ir);

    spl_ir2vm_init(&ir2vm, &a2ir.ir->ir, &lang_type.ir, &layout, &attr);
    ir2vm_ok = 1;
    spl_diag_set_stage(&diag, SPL_STAGE_IR2VM);
    ir2vm.diag = &diag;
    rc = spl_ir2vm_run(&ir2vm);
    if (spl_diag_has_error(&diag) || spl_diag_has_fatal(&diag)) {
        spl_diag_summary(&diag);
        goto cleanup;
    }
    if (dump == DUMP_VM) {
        spl_prog_dump(&ir2vm.prog);
        goto cleanup;
    }

    spl_prog_store_to_file(outpath, &ir2vm.prog);
    rc = 0;

cleanup:
    if (ir2vm_ok)
        spl_ir2vm_drop(&ir2vm);
    if (a2ir_ok)
        spl_ast2ir_drop(&a2ir);
    if (irb_ok)
        spl_ir_builder_drop(&ir_builder);
    if (sema_ok)
        spl_sema_drop(&sema);
    if (scope_ok)
        spl_scope_drop(&scope);
    if (layout_ok)
        spl_layout_drop(&layout);
    if (lang_type_ok)
        spl_lang_type_drop(&lang_type);
    if (attr_ok)
        spl_attr_drop(&attr);

    spl_ast_drop(&drv.main_ast);
    for (usize i = 0; i < vec_size(drv.modules); ++i) {
        spl_ast_t *a = vec_at(drv.modules, i).ast;
        if (a) {
            spl_ast_drop(a);
            free(a);
        }
        free(vec_at(drv.modules, i).abs_path);
        free(vec_at(drv.modules, i).src);
    }
    vec_free(drv.modules);
    /* map 的 key 是 module 里 strdup 出的 abs_path，已随 module 释放 */
    map_free(drv.path_to_idx);
    free(drv.main_path);
    spl_diag_drop(&diag);
    return rc;
}

int main(int argc, char **argv) {
    if (argc < 2) {
        LOG_FATAL("Usage: splc0 [--dump <flags>] <in> [out]");
        return 1;
    }
    if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-h") == 0) {
        LOG_INFO("splc0 <in> <out>              compile (.spl -> .sir, stage B)");
        LOG_INFO("splc0 --dump <flags> <file>   dump: lex,ast,sema,ir,vm,none");
        return 0;
    }
    int argi = 1;
    if (argi >= argc) {
        LOG_FATAL("Usage: splc0 [--dump <flags>] <in> [out]");
        return 1;
    }

    dump_t dump = DUMP_NONE;
    if (strcmp(argv[argi], "--dump") == 0) {
        if (argc < argi + 3) {
            LOG_INFO("splc0: --dump need <flags> <file>");
            return 1;
        }

        const char *flags = argv[argi + 1];
        if (strcmp(flags, "lex") == 0) {
            dump = DUMP_TOKEN;
        } else if (strcmp(flags, "ast") == 0) {
            dump = DUMP_AST;
        } else if (strcmp(flags, "sema") == 0) {
            dump = DUMP_SEMA;
        } else if (strcmp(flags, "ir") == 0) {
            dump = DUMP_IR;
        } else if (strcmp(flags, "vm") == 0) {
            dump = DUMP_VM;
        } else if (strcmp(flags, "none") == 0) {
            dump = DUMP_NONE;
        } else {
            LOG_FATAL("splc0: --dump takes one of lex|ast|sema|ir|vm|none, got `%s`", flags);
            return 1;
        }
        argi += 2;
    }
    /* splc0 [-g] <in> <out> */
    if (argc < argi + 2) {
        LOG_FATAL("Usage: splc0 <in> <out> [-g]  or  splc0 --dump <flags> <file>\n");
        return 1;
    }
    int gen_debug = 0;
    if (strcmp(argv[argi], "-g") == 0) {
        gen_debug = 1;
        argi++;
    }
    if (argc < argi + 2) {
        LOG_FATAL("Usage: splc0 <in> <out> [-g]  or  splc0 --dump <flags> <file>\n");
        return 1;
    }
    long len;
    char *src = read_file(argv[argi], &len);
    if (!src)
        return 1;
    int rc = compile_spl(src, argv[argi], argv[argi + 1], gen_debug, dump);
    free(src);
    return rc;
}