#include "asm/report.h"
#include "assemble.h"
#include "eval.h"
#include "parse.h"
#include "resolve.h"
#include "str.h"
#include <assert.h>
#include <errno.h>
#include <libgen.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    const char* input_file;
    const char* output_file;
} Args;

typedef struct {
    char* ident;
    Loc loc;
    uint16_t ip;
} GlobalSym;

typedef struct {
    size_t id;
    char* ident;
} ExternSym;

typedef struct {
    size_t id;
    uint16_t ip;
} ExternRef;

static inline Args parse_args(int argc, char** argv);
static inline char* read_text_file(const char* filename);
static inline int include_file(IdentResolver* resolver,
    OperandEvaluator* evaluator,
    const char* origin,
    const char* filename);
static inline void report_redefinition(
    const IdentResol* existing, PStmt* stmt, Reporter* rep);
static inline int define_labels(
    IdentResolver* resolver, PLabel* label, uint16_t asm_ip, Reporter* rep);
static inline void use_labels(IdentResolver* resolver, PLabel* label);

int main(int argc, char** argv)
{
    int res = 0;

    Args args = parse_args(argc, argv);

    char* input_text = read_text_file(args.input_file);

    bool errors_occured = false;

    Parser parser;
    parser_construct(&parser, args.input_file, input_text);

    Reporter rep = {
        .filename = parser.lexer.filename,
        .text = parser.lexer.text,
        .text_len = parser.lexer.text_len,
    };

    IdentResolver resolver;
    ident_resolver_construct(&resolver);

    OperandEvaluator evaluator;
    operand_evaluator_construct(&evaluator, &resolver, &rep);

    size_t lines_capacity = 1024;
    PLine** lines = malloc(sizeof(PLine*) * lines_capacity);
    size_t lines_size = 0;

    size_t global_symbols_capacity = 8;
    GlobalSym* global_symbols
        = malloc(sizeof(GlobalSym) * global_symbols_capacity);
    size_t global_symbols_size = 0;

    size_t extern_symbols_capacity = 8;
    ExternSym* extern_symbols
        = malloc(sizeof(ExternSym) * extern_symbols_capacity);
    size_t extern_symbols_size = 0;

    while (!parser_done(&parser)) {
        PStmt* stmt = parser_next_stmt(&parser);
        if (!stmt) {
            continue;
        }
        switch (stmt->ty) {
            case PStmtTy_Line:
                if (lines_size + 1 > lines_capacity) {
                    lines_capacity *= 2;
                    lines = realloc(lines, sizeof(PLine*) * lines_capacity);
                }
                lines[lines_size++] = stmt->line;
                // NOTE: Shallow free.
                free(stmt);
                break;
            case PStmtTy_Global:
                if (global_symbols_size + 1 > global_symbols_capacity) {
                    global_symbols_capacity *= 2;
                    global_symbols = realloc(global_symbols,
                        sizeof(GlobalSym) * global_symbols_capacity);
                }
                global_symbols[global_symbols_size++] = (GlobalSym) {
                    .ident = asm_strdup(stmt->ident),
                    .loc = stmt->loc,
                    .ip = 0,
                };
                pstmt_free(stmt);
                break;
            case PStmtTy_Extern: {
                const IdentResol* existing
                    = ident_resolver_resolve(&resolver, stmt->ident);
                if (existing != NULL) {
                    report_redefinition(existing, stmt, &rep);
                    pstmt_free(stmt);
                    continue;
                }
                size_t id = extern_symbols_size;
                ident_resolver_define_extern(
                    &resolver, asm_strdup(stmt->ident), stmt->loc, id);
                if (extern_symbols_size + 1 > extern_symbols_capacity) {
                    extern_symbols_capacity *= 2;
                    extern_symbols = realloc(extern_symbols,
                        sizeof(ExternSym) * extern_symbols_capacity);
                }
                extern_symbols[extern_symbols_size++] = (ExternSym) {
                    .ident = asm_strdup(stmt->ident),
                    .id = id,
                };
                pstmt_free(stmt);
                break;
            }
            case PStmtTy_Const: {
                const IdentResol* existing
                    = ident_resolver_resolve(&resolver, stmt->ident);
                if (existing != NULL) {
                    report_redefinition(existing, stmt, &rep);
                    pstmt_free(stmt);
                    continue;
                }
                EvaledOperand evaled
                    = eval_operand_to_imm(&evaluator, stmt->value);
                if (evaled.ty == EoTy_Err) {
                    pstmt_free(stmt);
                    errors_occured = true;
                    continue;
                }
                ident_resolver_define_const(&resolver,
                    asm_strdup(stmt->ident),
                    stmt->loc,
                    evaled.imm,
                    asm_strdup(args.input_file));
                pstmt_free(stmt);
                break;
            }
            case PStmtTy_Include: {
                int include_res = include_file(
                    &resolver, &evaluator, args.input_file, stmt->str);
                pstmt_free(stmt);
                if (include_res != 0) {
                    errors_occured = true;
                }
                break;
            }
        }
    }

    errors_occured &= parser_error_occured(&parser);

    size_t extern_refs_capacity = 8;
    ExternRef* extern_refs = malloc(sizeof(ExternRef) * extern_refs_capacity);
    size_t extern_refs_size = 0;

    size_t chunk_capacity = 64;
    uint16_t* chunk = malloc(sizeof(uint16_t) * chunk_capacity);

    uint16_t ip = 0;
    for (size_t i = 0; i < lines_size; ++i) {
        operand_evaluator_reset_externals(&evaluator);

        int res = define_labels(&resolver, lines[i]->labels, ip, &rep);
        if (res != 0)
            errors_occured = true;

        uint16_t size = pline_assemble(&evaluator, chunk, lines[i], &rep);
        if (size == 0)
            errors_occured = true;

        for (size_t i = 0; i < evaluator.externals_size; ++i) {
            const SurrogateExternal* ex = &evaluator.externals[i];
            bool found = false;
            for (uint16_t chunk_i = 1; chunk_i < size; ++chunk_i) {
                if (chunk[chunk_i] == ex->surrogate) {
                    found = true;

                    if (extern_refs_size + 1 > extern_refs_capacity) {
                        extern_refs_capacity *= 2;
                        extern_refs = realloc(extern_refs,
                            sizeof(ExternRef) * extern_refs_capacity);
                    }
                    extern_refs[extern_refs_size++] = (ExternRef) {
                        .id = ex->extern_id,
                        .ip = ip + chunk_i,
                    };

                    break;
                }
            }
            assert(found);
        }

        ip += size;
    }

    for (size_t i = 0; i < global_symbols_size; ++i) {
        const IdentResol* res
            = ident_resolver_resolve(&resolver, global_symbols[i].ident);
        if (res == NULL) {
            REPORTF_ERROR("undefined global '%s'", global_symbols[i].ident);
            REPORTF_INFO("'%s' declared global here", global_symbols[i].ident);
            reporter_print_loc(&rep, global_symbols[i].loc);
            continue;
        } else if (res->ty != IdentResolTy_Label) {
            REPORTF_ERROR("'%s' cannot be declared global", res->ident);
            reporter_print_loc(&rep, res->loc);
            REPORTF_INFO("'%s' declared global here", global_symbols[i].ident);
            reporter_print_loc(&rep, global_symbols[i].loc);
            continue;
        }
        global_symbols[i].ip = res->ip;
    }

    if (errors_occured) {
        fprintf(stderr, "nothing written. stopping...\n");
        res = -1;
        goto leave_free_chunk;
    }

    evaluator.second_pass = true;

    printf("globals:\n");
    printf("symbol          address\n");
    printf("-----------------------\n");
    for (size_t i = 0; i < global_symbols_size; ++i) {
        GlobalSym* sym = &global_symbols[i];
        printf("%-15s %d\n", sym->ident, sym->ip);
    }
    printf("\nexterns:\n");
    printf("id  symbol\n");
    printf("--------------------\n");
    for (size_t i = 0; i < extern_symbols_size; ++i) {
        ExternSym* sym = &extern_symbols[i];
        printf("%-3ld %s\n", sym->id, sym->ident);
    }
    printf("\nextern uses:\n");
    printf("id  address\n");
    printf("-----------\n");
    for (size_t i = 0; i < extern_refs_size; ++i) {
        ExternRef* ref = &extern_refs[i];
        printf("%3ld %04x\n", ref->id, ref->ip);
    }

    FILE* output_fp = fopen(args.output_file, "wb");
    if (!output_fp) {
        REPORTF_ERROR("could not open output file '%s': %s",
            args.output_file,
            strerror(errno));
        res = -1;
        goto leave_free_chunk;
    }

    size_t total_bytes_written = 0;
    for (size_t i = 0; i < lines_size; ++i) {
        use_labels(&resolver, lines[i]->labels);

        size_t size = pline_assemble(&evaluator, chunk, lines[i], &rep);
        if (size == 0) {
            errors_occured = true;
        }

        if (errors_occured)
            break;

        size_t bytes_written = fwrite(chunk, sizeof(uint16_t), size, output_fp);
        total_bytes_written += bytes_written;
        if (bytes_written != size) {
            REPORTF_ERROR("could not write to output file '%s': %s",
                args.output_file,
                strerror(errno));
            errors_occured = true;
        }
    }
    fclose(output_fp);

    if (errors_occured) {
        fprintf(
            stderr, "%ld bytes written. stopping...\n", total_bytes_written);
        res = -1;
        goto leave_free_chunk;
    }

    res = 0;
leave_free_chunk:
    free(chunk);
    // leave_free_lines:
    for (size_t i = 0; i < lines_size; ++i)
        pline_free(lines[i]);
    for (size_t i = 0; i < global_symbols_size; ++i)
        free(global_symbols[i].ident);
    free(global_symbols);
    for (size_t i = 0; i < extern_symbols_size; ++i)
        free(extern_symbols[i].ident);
    free(extern_symbols);
    free(extern_refs);
    free(lines);
    free(input_text);
    operand_evaluator_destroy(&evaluator);
    ident_resolver_destroy(&resolver);
    return res;
}

static inline Args parse_args(int argc, char** argv)
{
    const char* input_file = NULL;
    const char* output_file = NULL;
    for (int i = 1; i < argc; ++i) {
        if (strcmp(argv[i], "-o") == 0) {
            i += 1;
            if (i >= argc) {
                REPORTF_ERROR("%s", "no filename given to -o");
                exit(1);
            }
            output_file = argv[i];
        } else {
            if (input_file != NULL) {
                REPORTF_ERROR("%s", "multiple input files specified");
                exit(1);
            }
            input_file = argv[i];
        }
    }
    if (input_file == NULL) {
        REPORTF_ERROR("%s", "no input file");
        exit(1);
    }
    if (output_file == NULL) {
        output_file = "out.o";
    }
    return (Args) {
        input_file,
        output_file,
    };
}

static inline int include_file(IdentResolver* resolver,
    OperandEvaluator* evaluator,
    const char* origin_filename,
    const char* include_filename)
{

    char* origin_dir_buf = asm_strdup(origin_filename);
    const char* origin_dir = dirname(origin_dir_buf);

    char* filepath = calloc(
        strlen(origin_dir) + strlen(include_filename) + 2, sizeof(char));

    strcat(filepath, origin_dir);
    free(origin_dir_buf);

    strcat(filepath, "/");
    strcat(filepath, include_filename);

    char* text = read_text_file(filepath);
    if (!text)
        return -1;

    Parser parser;
    parser_construct(&parser, filepath, text);

    Reporter rep = {
        .filename = parser.lexer.filename,
        .text = parser.lexer.text,
        .text_len = parser.lexer.text_len,
    };

    bool errors_occured = false;

    while (!parser_done(&parser)) {
        PStmt* stmt = parser_next_stmt(&parser);
        if (!stmt) {
            continue;
        }
        switch (stmt->ty) {
            case PStmtTy_Const: {
                const IdentResol* existing
                    = ident_resolver_resolve(resolver, stmt->ident);
                if (existing != NULL) {
                    if (existing->ty == IdentResolTy_Const
                        && strcmp(existing->src_filename, filepath) == 0) {
                        // Same file included multiple times.
                        pstmt_free(stmt);
                        continue;
                    }
                    report_redefinition(existing, stmt, &rep);
                    errors_occured = true;
                    pstmt_free(stmt);
                    continue;
                }
                EvaledOperand evaled
                    = eval_operand_to_imm(evaluator, stmt->value);
                if (evaled.ty == EoTy_Err) {
                    pstmt_free(stmt);
                    errors_occured = true;
                    continue;
                }
                ident_resolver_define_const(resolver,
                    asm_strdup(stmt->ident),
                    stmt->loc,
                    evaled.imm,
                    asm_strdup(filepath));
                break;
            }
            case PStmtTy_Include: {
                int include_res
                    = include_file(resolver, evaluator, filepath, stmt->str);
                pstmt_free(stmt);
                if (include_res != 0) {
                    errors_occured = true;
                }
                break;
            }
            default:
                break;
        }
        pstmt_free(stmt);
    }

    errors_occured &= parser_error_occured(&parser);

    free(filepath);
    free(text);
    return errors_occured;
}

static inline void report_redefinition(
    const IdentResol* existing, PStmt* stmt, Reporter* rep)
{
    REPORTF_ERROR("redefinition of identifier '%s'", stmt->ident);
    reporter_print_loc(rep, stmt->loc);
    const char* filename = rep->filename;
    if (existing->ty == IdentResolTy_Const) {
        filename = existing->src_filename;
    }
    REPORTF_INFO("previous definition of '%s' here", existing->ident);
    reporter_print_loc(rep, existing->loc);
    rep->filename = filename;
}

static inline char* read_text_file(const char* filename)
{
    FILE* fp = fopen(filename, "r");
    if (!fp) {
        REPORTF_ERROR("could not open file '%s' for reading: %s",
            filename,
            strerror(errno));
        return NULL;
    }
    fseek(fp, 0L, SEEK_END);
    size_t file_size = (size_t)ftell(fp);
    rewind(fp);

    char* text = calloc(file_size + 1, sizeof(char));
    size_t bytes_read = fread(text, sizeof(char), file_size, fp);
    fclose(fp);
    if (bytes_read != file_size) {
        REPORTF_ERROR(
            "could not read input file '%s': %s", filename, strerror(errno));
        return NULL;
    }
    return text;
}

static inline int define_labels(
    IdentResolver* resolver, PLabel* label, uint16_t asm_ip, Reporter* rep)
{
    if (label == NULL)
        return 0;
    int res = define_labels(resolver, label->next, asm_ip, rep);
    const IdentResol* existing = ident_resolver_resolve(resolver, label->ident);
    if (existing != NULL) {
        REPORTF_ERROR("redefinition of identifier '%s'", label->ident);
        reporter_print_loc(rep, label->loc);

        const char* filename = rep->filename;
        if (existing->ty == IdentResolTy_Const) {
            rep->filename = existing->src_filename;
        }
        REPORTF_INFO("original definition of '%s'", existing->ident);
        reporter_print_loc(rep, existing->loc);
        rep->filename = filename;
        return 1;
    }
    if (label->sub_label) {
        ident_resolver_define_sublabel(
            resolver, asm_strdup(label->ident), label->loc, asm_ip);
    } else {
        ident_resolver_define_label(
            resolver, asm_strdup(label->ident), label->loc, asm_ip);
    }
    return res;
}

static inline void use_labels(IdentResolver* resolver, PLabel* label)
{
    if (label == NULL)
        return;
    use_labels(resolver, label->next);
    const IdentResol* existing = ident_resolver_resolve(resolver, label->ident);
    if (existing && existing->ty == IdentResolTy_Label) {
        resolver->current_parent = existing;
    }
}