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

#include "badlib.h"
#include "compiler.c"
#include "lexer.c"
#include "parser.c"
#include "semantic.c"
#include "vm.c"

// TODO: unions
// TODO: embed (binary file) and include (source file)
// TODO: revisit ast parsing for pointers and arrays (I think I'm missing corner
// cases).

typedef enum ExecMode {
    RUN_NORMAL,
    PRINT_LEX,
    PRINT_PARSE,
    PRINT_SEMANTIC,
    PRINT_SYMTABLES,
} ExecMode;

static ExecMode mode = RUN_NORMAL;

void
init(void) {
    log_init_default();
}

void
process_file(Str path) {
#if DEBUG == 1
    println("%s", path);
#endif
    Arena lexer_arena = arena_create(LEXER_MEM, os_allocator);

    FileContents file = platform_read_file(path, &lexer_arena);
    if (file.err) {
        eprintln("%s: error: %s", path, cstr("couldn't read the file"));
        exit(EXIT_FAILURE);
    }
    sz errors = 0;

    // Lexer.
    Scanner scanner = {.str = file.data};
    Token *tokens = NULL;
    Token tok = {0};
    while (tok.kind != TOK_EOF) {
        tok = scan_token(&scanner);
        if (tok.kind == TOK_UNKNOWN) {
            eprintln("%s:%d:%d:%s %s", path, tok.line, tok.col,
                     token_str[tok.kind], tok.val);
            errors++;
            continue;
        }
        array_push(tokens, tok, &lexer_arena);
    }
    // Only proceed if there are no errors.
    if (errors) {
        exit(EXIT_FAILURE);
    }
    if (mode == PRINT_LEX) {
        print_tokens(path, tokens);
        goto stop;
    }

    // Parser.
    Parser parser = {
        .tokens = tokens,
        .storage = &lexer_arena,
        .file_name = path,
    };
    array_init(parser.nodes, 256, parser.storage);
    parse_advance(&parser);
    while (parser.current.kind != TOK_EOF) {
#if DEBUG == 1
        static sz ctr = 0;
        println("ROOT: %d", ctr++);
#endif
        parse_expr(&parser, PREC_LOW);
        if (parser.panic) {
            break;
        }
    }
    parse_consume(&parser, TOK_EOF, cstr("expected end of file"));
    if (parser.err) {
        exit(EXIT_FAILURE);
    }
    if (mode == PRINT_PARSE) {
        graph_ast(parser.nodes);
        goto stop;
    }

    // Semantic analysis.
    Analyzer analyzer = (Analyzer){
        .storage = &lexer_arena,
        .file_name = path,
    };
    symbolic_analysis(&analyzer, &parser);
    if (analyzer.err) {
        exit(EXIT_FAILURE);
    }

    // Printing symbol tables.
    if (mode == PRINT_SYMTABLES) {
        graph_types(analyzer.scopes, lexer_arena);
    }
    if (mode == PRINT_SEMANTIC) {
        graph_ast(parser.nodes);
        goto stop;
    }

#if DEBUG == 1
    println("========== enums ==========");
    for (sz i = 0; i < array_size(analyzer.scopes); i++) {
        Arena scratch = lexer_arena;
        Scope *scope = analyzer.scopes[i];
        EnumMapIter iter = enummap_iterator(scope->enums, &scratch);
        EnumMap *m = enummap_next(&iter, &scratch);
        while (m) {
            println("scope: %x{2} -- %s: enum: %s", scope->id, path,
                    m->val.name);
            m = enummap_next(&iter, &scratch);
        }
    }
    println("========= structs =========");
    for (sz i = 0; i < array_size(analyzer.scopes); i++) {
        Arena scratch = lexer_arena;
        Scope *scope = analyzer.scopes[i];
        StructMapIter iter = structmap_iterator(scope->structs, &scratch);
        StructMap *m = structmap_next(&iter, &scratch);
        while (m) {
            println("scope: %x{2} -- %s: struct: %s", scope->id, path,
                    m->val.name);
            m = structmap_next(&iter, &scratch);
        }
    }
    println("======== functions ========");
    for (sz i = 0; i < array_size(analyzer.scopes); i++) {
        Arena scratch = lexer_arena;
        Scope *scope = analyzer.scopes[i];
        FunMapIter iter = funmap_iterator(scope->funcs, &scratch);
        FunMap *m = funmap_next(&iter, &scratch);
        while (m) {
            println("scope: %x{2} -- %s: func: %s(%s): (%s)", scope->id, path,
                    m->val.name, m->val.param_type, m->val.return_type);
            m = funmap_next(&iter, &scratch);
        }
    }
    println("========= symbols =========");
    for (sz i = 0; i < array_size(analyzer.scopes); i++) {
        Arena scratch = lexer_arena;
        Scope *scope = analyzer.scopes[i];
        SymbolMapIter iter = symmap_iterator(scope->symbols, &scratch);
        SymbolMap *m = symmap_next(&iter, &scratch);
        while (m) {
            println("scope: %x{2} -- %s: %s %s: %s", scope->id, path,
                    sym_kind_str[m->val.kind], m->key, m->val.name);
            m = symmap_next(&iter, &scratch);
        }
    }
#endif

    // TODO: Type checking.

    // Compile roots.
    Arena bytecode_arena = arena_create(LEXER_MEM, os_allocator);
    Chunk chunk = {.file_name = path, .storage = &bytecode_arena};
    array_zero(chunk.constants, 256, &bytecode_arena);
    array_zero(chunk.code, 0xffff, &bytecode_arena);
    sz n_roots = array_size(parser.nodes);
    CompResult res = {0};
    for (sz i = 0; i < n_roots; i++) {
        // The parser stores the root nodes as a stack.
        Node *root = parser.nodes[i];
        res = compile_expr(&chunk, root);
        if (res.type == COMP_ERR) {
            eprintln("compilation error...");
            exit(EXIT_FAILURE);
        }
    }
    sz res_reg = 0;
    switch (res.type) {
        case COMP_CONST: {
            res_reg = chunk.reg_idx++;
            Instruction inst =
                (Instruction){.op = OP_LD64K, .dst = res_reg, .a = res.idx};
            array_push(chunk.code, inst, chunk.storage);
        } break;
        case COMP_REG: {
            res_reg = res.idx;
        } break;
        default: break;
    }
    // After we are done move the last result to r0 for printing.
    Instruction halt = (Instruction){.op = OP_HALT, .dst = res_reg};
    array_push(chunk.code, halt, &bytecode_arena);

    if (chunk.const_idx >= 256) {
        eprintln("too many constants on chunk %s", chunk.id);
        exit(EXIT_FAILURE);
    }
    if (chunk.str_idx >= 256) {
        eprintln("too many strings on chunk %s", chunk.id);
        exit(EXIT_FAILURE);
    }
    if (chunk.reg_idx >= 256) {
        eprintln("too many registers used on chunk %s", chunk.id);
        exit(EXIT_FAILURE);
    }

    disassemble_chunk(chunk);

    // Run bytecode on VM.
    VM vm = {0};
    vm_init(&vm, &chunk);
    // println("VM REGISTERS BEFORE:\n%{Mem}",
    //         &(Array){.mem = (u8 *)&vm.regs, sizeof(vm.regs)});
    vm_run(&vm);
    // println("VM REGISTERS AFTER:\n%{Mem}",
    //         &(Array){.mem = (u8 *)&vm.regs, sizeof(vm.regs)});
    // println("VM MEMORY AFTER:\n%{Mem}",
    //         &(Array){.mem = (u8 *)&vm.stack, sizeof(vm.stack)});

#if DEBUG == 1
    println("Space used: %{Arena}", &lexer_arena);
#endif

    // arena_destroy(&bytecode_arena, os_allocator);
stop:
    // Free up resources.
    arena_destroy(&lexer_arena, os_allocator);
}

#ifndef BIN_NAME
#define BIN_NAME "bdl"
#endif

void
print_usage(void) {
    printf("Usage: %s [options] <filename filename ...>\n", BIN_NAME);
    printf("\n");
    printf("\t-h \t\tShow usage.\n");
    printf(
        "\t-p [l|p|s|t]\tPrint mode for [l]exing, [p]arsing, "
        "[s]emantic "
        "analysis, symbol [t]ables\n");
    printf("\n");
}

int
main(int argc, char *argv[]) {
    int option;
    while ((option = getopt(argc, argv, "hp:")) != -1) {
        switch (option) {
            case 'h': {
                print_usage();
                goto exit_success;
            } break;
            case 'p': {
                if (optarg[0] == 'l' && optarg[1] == '\0') {
                    mode = PRINT_LEX;
                } else if (optarg[0] == 'p' && optarg[1] == '\0') {
                    mode = PRINT_PARSE;
                } else if (optarg[0] == 's' && optarg[1] == '\0') {
                    mode = PRINT_SEMANTIC;
                } else if (optarg[0] == 't' && optarg[1] == '\0') {
                    mode = PRINT_SYMTABLES;
                } else {
                    print_usage();
                    return EXIT_FAILURE;
                }
            } break;
            default: {
                print_usage();
                return EXIT_FAILURE;
            } break;
        }
    }

    init();

    // Run from stdin.
    if (optind == argc) {
        // TODO: REPL
        // repl();
        goto exit_success;
    }

    // Run from file.
    while (optind < argc) {
        char *file_name = argv[optind];
        Str file_path = STR(file_name);
        process_file(file_path);
        optind++;
    }

exit_success:
    return EXIT_SUCCESS;
}