#define __SCC_CMD_LOG_IMPL__
#include "cmd_log.h"

#include <argparse.h>
#include <scc_lexer.h>
#include <scc_parser.h>
#include <scc_pproc.h>

#include <scc_ast2ir.h>
#include <scc_ast_dump.h>
#include <scc_hir2lir.h>
#include <scc_hir_dump.h>
#include <scc_lir2mir.h>
#include <scc_lir_dump.h>
#include <scc_mir_dump.h>

#include <scc_mir_pass.h>

#include <scc_ir2mcode.h>
#include <scc_ir2sccf.h>
#include <sccf2pe.h>
#include <sccf_linker.h>

#define __SCC_TYPE_ABI_WIN_X64_IMPL__
#include <scc_type_abi_win_x64.h>

#include "config.h"
void init_platform(void);
#define GET_VALID_FP(fp) (fp == nullptr ? scc_stdout : fp)

static void print_ring(scc_lexer_tok_ring_t *ring, int verbose) {
    scc_lexer_tok_t tok = {0};
    int ret = 0;
    while (1) {
        scc_ring_next_consume(*ring, tok, ret);
        if (ret == false || tok.type == SCC_TOK_EOF) {
            break;
        }
        if (verbose == 0) {
            scc_printf("%s ", scc_get_tok_name(tok.type));
        } else if (verbose >= 1) {
            scc_printf(
                "token [%-8s] `%s` at %s:%d:%d\n", scc_get_tok_name(tok.type),
                tok.type != SCC_TOK_ENDLINE ? scc_str_as_cstr(&tok.lexeme)
                                            : "\\n",
                tok.loc.name, tok.loc.line, tok.loc.col);
        }
        scc_lexer_tok_drop(&tok);
    }
}

static void print_file(scc_lexer_tok_ring_t *ring, scc_file_t fp) {
    scc_lexer_tok_t tok = {0};
    int ret = 0;

    while (1) {
        scc_ring_next_consume(*ring, tok, ret);
        if (ret == false || tok.type == SCC_TOK_EOF) {
            break;
        }
        if (fp == scc_stdout) {
            scc_printf("%s", scc_str_as_cstr(&tok.lexeme));
        } else {
            usize ret = scc_fwrite(fp, scc_str_as_cstr(&tok.lexeme),
                                   scc_str_len(&tok.lexeme));
            if (ret != scc_str_len(&tok.lexeme)) {
                LOG_FATAL("Failed to write to file");
            }
        }
        scc_lexer_tok_drop(&tok);
    }
    scc_fclose(fp);
}

static void tree_dump_output(const char *str, usize len, void *user) {
    scc_fprintf(user, "%.*s", (int)len, str);
}

static inline int mir_dump(scc_file_t fp, scc_mir_module_t *mir_module) {
    scc_tree_dump_t tree_dump;
    scc_mir_dump_ctx_t mir_dump_ctx;
    if (fp == nullptr) {
        scc_tree_dump_init(&tree_dump, true);
    } else {
        scc_tree_dump_init(&tree_dump, false);
    }
    scc_mir_dump_init(&mir_dump_ctx, &tree_dump, mir_module);
    scc_mir_dump_module(&mir_dump_ctx);

    scc_tree_dump_flush(&tree_dump, tree_dump_output, GET_VALID_FP(fp));
    scc_tree_dump_drop(&tree_dump);
    return 0;
}

// Compile a single source file to an SCCF builder.
// Returns 0 on success, non-zero on error.
// The caller must keep the builder alive until after linking.
static int compile_file_to_sccf(const char *input_file, scc_config_t *config,
                                sccf_builder_t *builder) {
    int error_code = 0;
    scc_sstream_t sstream;
    error_code = scc_sstream_init(&sstream, input_file, 1024);
    if (error_code) {
        return error_code;
    }

    LOG_INFO("=== phase 1/6: lexing '%s' ===", input_file);
    scc_lexer_t lexer;
    scc_lexer_init(&lexer, scc_sstream_to_ring(&sstream));
    if (config->emit_stage == SCC_EMIT_STAGE_LEX) {
        // Early-lex is handled per-file at call site
        scc_lexer_drop(&lexer);
        scc_sstream_drop(&sstream);
        return -1;
    }

    LOG_INFO("=== phase 2/6: preprocessing '%s' ===", input_file);
    scc_pproc_t pproc;
    scc_pproc_init(&pproc, scc_lexer_to_ring(&lexer, 8, true));
    scc_pproc_add_include_path_cstr(&pproc, "./.scc_include");
    scc_vec_foreach(config->include_paths, i) {
        scc_pproc_add_include_path_cstr(&pproc,
                                        scc_vec_at(config->include_paths, i));
    }
    scc_lexer_tok_vec_t pproc_tok_vec;
    scc_vec_init(pproc_tok_vec);
    scc_lexer_tok_t tok = {
        .lexeme = scc_str_from_cstr("1"),
        .type = SCC_TOK_INT_LITERAL,
        .loc.name = "<internal>",
        .loc.line = 0,
        .loc.col = 0,
        .loc.offset = 0,
    };
    scc_vec_push(pproc_tok_vec, tok);
    scc_str_t pproc_predefined_macros[] = {
        scc_str_from_cstr("__SCC__"),
        scc_str_from_cstr("_WIN64"),
        scc_str_from_cstr("_WIN32"),
        scc_str_from_cstr("__x86_64__"),
    };
    for (usize i = 0; i < SCC_ARRLEN(pproc_predefined_macros); i += 1) {
        scc_vec_init(pproc_tok_vec);
        scc_lexer_tok_t coped_tok = scc_lexer_tok_copy(&tok);
        scc_vec_push(pproc_tok_vec, coped_tok);
        scc_pproc_add_object_macro(&pproc.macro_table,
                                   &pproc_predefined_macros[i], &pproc_tok_vec);
    }

    scc_lexer_tok_ring_t *tok_ring =
        scc_pproc_to_ring(&pproc, 1024, false, false);
    LOG_INFO("=== phase 3/6: parsing '%s' ===", input_file);
    scc_parser_t parser;
    scc_ast_module_t ast_module;
    scc_ast_module_init(&ast_module);
    scc_sema_ctx_t sema_ctx;
    scc_sema_init(&sema_ctx, &ast_module);
    scc_parser_init(&parser, tok_ring, &ast_module, &sema_ctx);
    scc_ast_translation_unit_t *translation_unit =
        scc_parse_translation_unit(&parser);
    if (parser.errcode != 0) {
        error_code = parser.errcode;
        goto cleanup_pproc;
    }

    LOG_INFO("=== phase 4/6: AST -> HIR ===");
    scc_ast2ir_ctx_t ast2ir_ctx;
    scc_hir_cprog_t cprog;
    scc_hir_cprog_init(&cprog, &SCC_TYPE_ABI_WIN_X64);
    scc_ast2ir_ctx_init(&ast2ir_ctx, &SCC_TYPE_ABI_WIN_X64, &ast_module,
                        &cprog);
    scc_ast2ir_run(&ast2ir_ctx, translation_unit);
    scc_ast2ir_ctx_drop(&ast2ir_ctx);

    LOG_INFO("=== phase 5/6: HIR -> LIR -> MIR ===");
    // HIR -> LIR
    scc_lir_module_t lir_module;
    scc_lir_module_init(&lir_module);
    scc_hir2lir(&lir_module, &cprog);

    // LIR -> MIR
    scc_mir_module_t mir_module;
    scc_mir_module_init(&mir_module);
    scc_lir2mir(&mir_module, &lir_module);

    // MIR passes
    scc_mir_pass(&mir_module, SCC_MIR_STAGE_ANY);

    LOG_INFO("=== phase 6/6: MIR -> SCCF ===");
    // MIR -> SCCF
    // NOTE: scc_ir2sccf calls sccf_builder_init(builder) internally
    scc_ir2sccf(builder, &mir_module);

    // Cleanup per-file resources (reverse order)
    scc_mir_module_drop(&mir_module);
    scc_lir_module_drop(&lir_module);
    scc_hir_cprog_drop(&cprog);
    scc_sema_drop(&sema_ctx);
    scc_parser_drop(&parser);
    scc_pproc_drop(&pproc);
    scc_lexer_drop(&lexer);
    scc_sstream_drop(&sstream);

    return 0;

cleanup_pproc:
    scc_pproc_drop(&pproc);
    scc_lexer_drop(&lexer);
    scc_sstream_drop(&sstream);
    return error_code;
}

int main(int argc, const char **argv, const char **envp) {
    init_platform();

#ifndef SCC_DEFAULT_ARGPARSE_LANG
#define SCC_DEFAULT_ARGPARSE_LANG SCC_ARGPARSE_LANG_ZH
#endif

    scc_argparse_lang_t argparse_lang = SCC_DEFAULT_ARGPARSE_LANG;
    for (const char **env = envp; *env != nullptr; env++) {
        const char *env_str = *env;
        if (scc_strcmp(env_str, "LANG=zh_CN.UTF-8") == 0) {
            argparse_lang = SCC_ARGPARSE_LANG_ZH;
        }
    }

    scc_config_t config = {
        .verbose = 0,
        .output_file = nullptr,
        .entry_point_symbol = nullptr,
        .emit_stage = SCC_EMIT_STAGE_DEFAULT,
        .target_description = "x86_64-pc-windows-msvc",
    };
    scc_vec_init(config.input_files);
    scc_vec_init(config.include_paths);
    scc_vec_init(config.define_macros);
    scc_vec_init(config.log_configs);

    scc_argparse_t argparse;
    setup_argparse(&argparse, &config, argparse_lang);
    int ret = scc_argparse_parse(&argparse, argc, argv);
    if (ret != 0) {
        scc_argparse_drop(&argparse);
        return 0;
    }
    scc_argparse_drop(&argparse);

    // // Apply log configuration from --log and -v
    // scc_log_set_verbosity(config.verbose);
    // scc_vec_foreach(config.log_configs, i) {
    //     scc_log_apply_config(scc_vec_at(config.log_configs, i));
    // }

    usize num_inputs = scc_vec_size(config.input_files);
    if (num_inputs == 0) {
        LOG_ERROR("no input files");
        return 1;
    }

    scc_file_t fp = nullptr;
    if (config.output_file) {
        cbool is_stdout = scc_strcmp(config.output_file, "-") == 0;
        if (!is_stdout) {
            fp = scc_fopen(config.output_file, SCC_FILE_WRITE);
            if (fp == nullptr) {
                LOG_FATAL("Failed to open file %s", config.output_file);
                return 1;
            }
        } else {
            fp = scc_stdout;
        }
    }

    // Single-file mode: original fast path
    if (num_inputs == 1) {
        const char *input_file = scc_vec_at(config.input_files, 0);

        scc_sstream_t sstream;
        int error_code = scc_sstream_init(&sstream, input_file, 1024);
        if (error_code) {
            return error_code;
        }

        scc_lexer_t lexer;
        scc_lexer_init(&lexer, scc_sstream_to_ring(&sstream));
        if (config.emit_stage == SCC_EMIT_STAGE_LEX) {
            scc_lexer_tok_ring_t *tok_ring =
                scc_lexer_to_ring(&lexer, 8, fp == nullptr ? false : true);
            if (fp == nullptr) {
                print_ring(tok_ring, config.verbose);
            } else {
                print_file(tok_ring, fp);
            }
            scc_lexer_drop_ring(tok_ring);
            scc_lexer_drop(&lexer);
            scc_sstream_drop(&sstream);
            return 0;
        }

        scc_pproc_t pproc;
        scc_pproc_init(&pproc, scc_lexer_to_ring(&lexer, 8, true));
        scc_pproc_add_include_path_cstr(&pproc, "./.scc_include");
        scc_vec_foreach(config.include_paths, i) {
            scc_pproc_add_include_path_cstr(
                &pproc, scc_vec_at(config.include_paths, i));
        }
        scc_lexer_tok_vec_t pproc_tok_vec;
        scc_vec_init(pproc_tok_vec);
        scc_lexer_tok_t tok = {
            .lexeme = scc_str_from_cstr("1"),
            .type = SCC_TOK_INT_LITERAL,
            .loc.name = "<internal>",
            .loc.line = 0,
            .loc.col = 0,
            .loc.offset = 0,
        };
        scc_vec_push(pproc_tok_vec, tok);
        scc_str_t pproc_predefined_macros[] = {
            scc_str_from_cstr("__SCC__"),
            scc_str_from_cstr("_WIN64"),
            scc_str_from_cstr("_WIN32"),
            scc_str_from_cstr("__x86_64__"),
        };
        for (usize i = 0; i < SCC_ARRLEN(pproc_predefined_macros); i += 1) {
            scc_vec_init(pproc_tok_vec);
            scc_lexer_tok_t coped_tok = scc_lexer_tok_copy(&tok);
            scc_vec_push(pproc_tok_vec, coped_tok);
            scc_pproc_add_object_macro(&pproc.macro_table,
                                       &pproc_predefined_macros[i],
                                       &pproc_tok_vec);
        }
        if (config.emit_stage == SCC_EMIT_STAGE_PP) {
            scc_lexer_tok_ring_t *tok_ring =
                scc_pproc_to_ring(&pproc, 8, true, true);
            if (fp == nullptr) {
                print_ring(tok_ring, config.verbose);
            } else {
                print_file(tok_ring, fp);
            }
            scc_pproc_drop(&pproc);
            scc_lexer_drop(&lexer);
            scc_sstream_drop(&sstream);
            return 0;
        }

        scc_lexer_tok_ring_t *tok_ring =
            scc_pproc_to_ring(&pproc, 1024, false, false);
        scc_parser_t parser;
        scc_ast_module_t ast_module;
        scc_ast_module_init(&ast_module);
        scc_sema_ctx_t sema_ctx;
        scc_sema_init(&sema_ctx, &ast_module);
        scc_parser_init(&parser, tok_ring, &ast_module, &sema_ctx);
        scc_ast_translation_unit_t *translation_unit =
            scc_parse_translation_unit(&parser);
        if (parser.errcode != 0) {
            scc_sema_drop(&sema_ctx);
            scc_parser_drop(&parser);
            scc_pproc_drop(&pproc);
            scc_lexer_drop(&lexer);
            scc_sstream_drop(&sstream);
            return parser.errcode;
        }

        scc_sema_drop(&sema_ctx);
        scc_parser_drop(&parser);
        scc_pproc_drop(&pproc);
        scc_lexer_drop(&lexer);
        scc_sstream_drop(&sstream);

        if (config.emit_stage == SCC_EMIT_STAGE_AST) {
            scc_tree_dump_t tree_dump;
            if (fp == nullptr) {
                scc_tree_dump_init(&tree_dump, true);
            } else {
                scc_tree_dump_init(&tree_dump, false);
            }
            scc_ast_dump_node(&tree_dump, (scc_ast_node_t *)translation_unit);
            scc_tree_dump_flush(&tree_dump, tree_dump_output, GET_VALID_FP(fp));
            scc_tree_dump_drop(&tree_dump);
            return 0;
        }

        scc_ast2ir_ctx_t ast2ir_ctx;
        scc_hir_cprog_t cprog;
        scc_hir_cprog_init(&cprog, &SCC_TYPE_ABI_WIN_X64);
        scc_ast2ir_ctx_init(&ast2ir_ctx, &SCC_TYPE_ABI_WIN_X64, &ast_module,
                            &cprog);
        scc_ast2ir_run(&ast2ir_ctx, translation_unit);
        scc_ast2ir_ctx_drop(&ast2ir_ctx);

        if (config.emit_stage == SCC_EMIT_STAGE_HIR) {
            scc_hir_dump_t ir_dump_ctx;
            scc_tree_dump_t tree_dump;
            if (fp == nullptr) {
                scc_tree_dump_init(&tree_dump, true);
            } else {
                scc_tree_dump_init(&tree_dump, false);
            }
            scc_hir_dump_init(&ir_dump_ctx, &tree_dump, &cprog);
            scc_hir_dump_cprog_linear(&ir_dump_ctx);

            scc_tree_dump_flush(&tree_dump, tree_dump_output, GET_VALID_FP(fp));
            scc_tree_dump_drop(&tree_dump);
            return 0;
        }

        scc_lir_module_t lir_module;
        scc_lir_module_init(&lir_module);
        scc_hir2lir(&lir_module, &cprog);
        if (config.emit_stage == SCC_EMIT_STAGE_LIR) {
            scc_lir_dump_ctx_t lir_dump_ctx;
            scc_tree_dump_t tree_dump;
            if (fp == nullptr) {
                scc_tree_dump_init(&tree_dump, true);
            } else {
                scc_tree_dump_init(&tree_dump, false);
            }
            scc_lir_dump_init(&lir_dump_ctx, &tree_dump, &lir_module);
            scc_lir_dump_module(&lir_dump_ctx);

            scc_tree_dump_flush(&tree_dump, tree_dump_output, GET_VALID_FP(fp));
            scc_tree_dump_drop(&tree_dump);
            return 0;
        }

        scc_mir_module_t mir_module;
        scc_mir_module_init(&mir_module);
        scc_lir2mir(&mir_module, &lir_module);

        switch (config.emit_stage) {
        case SCC_EMIT_STAGE_MIR:
            return mir_dump(fp, &mir_module);
        case SCC_EMIT_STAGE_MIR_PASS_REG_ALLOC:
            scc_mir_pass(&mir_module, SCC_MIR_STAGE_REGALLOC);
            return mir_dump(fp, &mir_module);
        case SCC_EMIT_STAGE_MIR_PASS_FLAME_LAYOUT:
            scc_mir_pass(&mir_module, SCC_MIR_STAGE_FRAME_LAYOUT);
            return mir_dump(fp, &mir_module);
        case SCC_EMIT_STAGE_MIR_PASS_PROLOG_EPILOG:
            scc_mir_pass(&mir_module, SCC_MIR_STAGE_PROLOGUE_EPILOGUE);
            return mir_dump(fp, &mir_module);
        default:
            scc_mir_pass(&mir_module, SCC_MIR_STAGE_ANY);
            break;
        }

        if (config.emit_stage == SCC_EMIT_STAGE_FLATBIN) {
            scc_mcode_t mcode = {0};
            scc_mcode_init(&mcode, SCC_MCODE_ARCH_X86_64);
            scc_ir2mcode(&mcode, nullptr, &mir_module);
            if (fp == nullptr) {
                LOG_WARN("emit flatbin can't write to stdout");
                return 0;
            }
            usize mcode_size = scc_mcode_size(&mcode);
            usize fwritten =
                scc_fwrite(fp, scc_mcode_unsafe_data(&mcode), mcode_size);
            if (fwritten != mcode_size) {
                LOG_ERROR("write flatbin failed, write %zu but need %zu",
                          fwritten, mcode_size);
                return 1;
            }
            return 0;
        }

        // Single-file: compile to SCCF then directly to PE
        sccf_builder_t sccf_builder = {0};
        scc_ir2sccf(&sccf_builder, &mir_module);
        sccf_builder_set_entry_symbol_name(&sccf_builder,
                                           config.entry_point_symbol);
        const sccf_t *sccf = sccf_builder_to_sccf(&sccf_builder);
        if (config.emit_stage == SCC_EMIT_STAGE_SCCF) {
            if (fp == nullptr) {
                scc_printf("output exe at %s\n", config.output_file);
            } else {
                sccf_builder_to_file(&sccf_builder, config.output_file);
            }
            return 0;
        }

        if (config.emit_stage == SCC_EMIT_STAGE_DEFAULT ||
            config.emit_stage == SCC_EMIT_STAGE_TARGET) {
            scc_pe_builder_t pe_builder;
            sccf2pe(&pe_builder, sccf);
            if (fp == nullptr) {
                scc_printf("output exe at %s\n", config.output_file);
            } else {
                scc_pe_dump_to_file(&pe_builder, config.output_file);
            }
            return 0;
        }

        Panic("unknown emit stage");
        return 1;
    }

    // Multi-file mode (num_inputs > 1)
    if (config.emit_stage != SCC_EMIT_STAGE_DEFAULT) {
        LOG_ERROR("--emit is not supported with multiple input files");
        return 1;
    }

    // First pass: compile each file to SCCF
    sccf_linker_t linker;
    sccf_linker_init(&linker);
    sccf_linker_set_entry_symbol_name(&linker, config.entry_point_symbol);

    scc_vec_foreach(config.input_files, i) {
        const char *input = scc_vec_at(config.input_files, i);
        sccf_builder_t builder = {0};

        int err = compile_file_to_sccf(input, &config, &builder);
        if (err != 0) {
            return err;
        }

        // Transfer this file's SCCF to the linker.
        // The builder data stays alive on the heap; the linker stores a shallow
        // copy. We must NOT drop/free the builder while the linker references
        // it.
        sccf_linker_add_sccf(&linker, sccf_builder_to_sccf(&builder));
    }

    // Second pass: link all SCCF modules and emit PE
    const sccf_t *linked = sccf_linker_link(&linker);
    scc_pe_builder_t pe_builder;
    sccf2pe(&pe_builder, linked);
    if (fp == nullptr) {
        scc_printf("output exe at %s\n", config.output_file);
    } else {
        scc_pe_dump_to_file(&pe_builder, config.output_file);
    }

    return 0;
}