#include "cf.h"
#include "da.h"
#include "dc.h"
#include "js.h"
#include "mem.h"
#include "p.h"

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

#define TOBI_VERSION "0.1.0"

static void usage(FILE *f) {
    fputs("usage: tobi [opts] <file> [file...]\n"
          "\n"
          "options:\n"
          "  --dis       print opcode-level disassembly\n"
          "  --json      print parsed ir as json\n"
          "  --dot       print recovered control flow as Graphviz DOT\n"
          "  --raw       print input/table metadata and aml offsets\n"
          "  --strict    treat malformed or unsupported constructs as hard errors\n"
          "  --plain     disable coloured diagnostics\n"
          "  --help      print usage\n"
          "  --version   print version\n", f);
}

static int read_file(const char *path, uint8_t **data, size_t *len) {
    FILE *fp = fopen(path, "rb");
    if (!fp) {
        fprintf(stderr, "tobi: %s: %s\n", path, strerror(errno));
        return 0;
    }
    if (fseek(fp, 0, SEEK_END) != 0) {
        fprintf(stderr, "tobi: %s: seek failed: %s\n", path, strerror(errno));
        fclose(fp);
        return 0;
    }
    long n = ftell(fp);
    if (n < 0) {
        fprintf(stderr, "tobi: %s: size failed: %s\n", path, strerror(errno));
        fclose(fp);
        return 0;
    }
    rewind(fp);
    *data = tobi_xmalloc((size_t)n ? (size_t)n : 1);
    *len = (size_t)n;
    if (*len && fread(*data, 1, *len, fp) != *len) {
        fprintf(stderr, "tobi: %s: read failed%s%s\n", path, ferror(fp) ? ": " : "",
                ferror(fp) ? strerror(errno) : "");
        fclose(fp);
        free(*data);
        *data = NULL;
        *len = 0;
        return 0;
    }
    if (fclose(fp) != 0) {
        fprintf(stderr, "tobi: %s: close failed: %s\n", path, strerror(errno));
        free(*data);
        *data = NULL;
        *len = 0;
        return 0;
    }
    return 1;
}

static void print_diags(const tobi_diag_list *dl, int plain) {
    int colour = !plain && isatty(fileno(stderr));
    for (size_t i = 0; i < dl->len; i++) {
        const char *lvl = tobi_diag_level_name(dl->items[i].level);
        const char *c = dl->items[i].level == TOBI_DIAG_ERROR ? "\033[31m" : "\033[33m";
        if (dl->items[i].has_source && dl->items[i].source) {
            fprintf(stderr, "%s:", dl->items[i].source);
        }
        if (colour) {
            fprintf(stderr, "0x%zx: %s%s\033[0m: %s\n", dl->items[i].off, c, lvl, dl->items[i].msg);
        } else {
            fprintf(stderr, "0x%zx: %s: %s\n", dl->items[i].off, lvl, dl->items[i].msg);
        }
    }
}

static void set_source(tobi_input_meta *meta, const char *path) {
    if (!path) {
        return;
    }
    (void)snprintf(meta->source, sizeof(meta->source), "%s", path);
}

static void apply_source(tobi_parse_result *res, const char *path, size_t input_index) {
    set_source(&res->meta, path);
    if (input_index < res->input_len) {
        set_source(&res->inputs[input_index], path);
    }
    tobi_diag_set_source(&res->diag, path, input_index);
    tobi_ns_set_source(&res->ns, path, input_index);
    if (res->root) {
        tobi_ir_set_source_recursive(res->root, path, input_index);
    }
}

static void print_raw_meta(const char *input, const tobi_parse_result *res) {
    printf("input=%s\nsource=%s\nformat=%s\nis_table=%s\ntable_length=%zu\naml_offset=%zu\naml_length=%zu\nrevision=%u\nchecksum_byte=0x%02x\nchecksum_sum=0x%02x\nchecksum_valid=%s\noem_id=%s\noem_table_id=%s\noem_revision=%u\ncreator_id=%s\ncreator_revision=%u\nheader_hash=0x%08x\naml_hash=0x%08x\nnamespace_objects=%zu\n",
           input, res->meta.source, res->meta.signature, res->meta.is_table ? "true" : "false",
           res->meta.table_len, res->meta.aml_off, res->meta.aml_len,
           (unsigned)res->meta.revision, (unsigned)res->meta.checksum_byte,
           (unsigned)res->meta.checksum_sum, res->meta.checksum_valid ? "true" : "false",
           res->meta.oem_id, res->meta.oem_table_id, res->meta.oem_revision,
           res->meta.creator_id, res->meta.creator_revision, res->meta.header_hash,
           res->meta.aml_hash, res->ns.len);
}

int main(int argc, char **argv) {
    int mode_dis = 0;
    int mode_json = 0;
    int mode_dot = 0;
    int mode_raw = 0;
    int strict = 0;
    int plain = 0;
    const char **files = NULL;
    size_t file_len = 0;
    size_t file_cap = 0;
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--help") == 0) {
            usage(stdout);
            return 0;
        }
        if (strcmp(argv[i], "--version") == 0) {
            puts("tobi " TOBI_VERSION);
            return 0;
        }
        if (strcmp(argv[i], "--dis") == 0) {
            mode_dis = 1;
        } else if (strcmp(argv[i], "--json") == 0) {
            mode_json = 1;
        } else if (strcmp(argv[i], "--dot") == 0) {
            mode_dot = 1;
        } else if (strcmp(argv[i], "--raw") == 0) {
            mode_raw = 1;
        } else if (strcmp(argv[i], "--strict") == 0) {
            strict = 1;
        } else if (strcmp(argv[i], "--plain") == 0) {
            plain = 1;
        } else if (argv[i][0] == '-') {
            fprintf(stderr, "tobi: unknown option %s\n", argv[i]);
            usage(stderr);
            return 2;
        } else {
            if (file_len == file_cap) {
                file_cap = tobi_xgrow_cap(file_cap, file_len + 1u, 4u);
                files = tobi_xrealloc(files, tobi_xmul_size(file_cap, sizeof(files[0])));
            }
            files[file_len++] = argv[i];
        }
    }
    if (file_len == 0) {
        usage(stderr);
        free(files);
        return 2;
    }
    if ((mode_dis + mode_json + mode_dot + mode_raw) > 1) {
        fprintf(stderr, "tobi: choose only one output mode\n");
        free(files);
        return 2;
    }
    uint8_t **data = tobi_xcalloc(file_len, sizeof(data[0]));
    size_t *lens = tobi_xcalloc(file_len, sizeof(lens[0]));
    for (size_t i = 0; i < file_len; i++) {
        if (!read_file(files[i], &data[i], &lens[i])) {
            for (size_t j = 0; j < file_len; j++) {
                free(data[j]);
            }
            free(lens);
            free(data);
            free(files);
            return 1;
        }
    }
    if (mode_dis || mode_raw) {
        int fail_any = 0;
        for (size_t i = 0; i < file_len; i++) {
            tobi_parse_result one;
            int ok = tobi_parse(data[i], lens[i], strict, &one);
            apply_source(&one, files[i], 0);
            if (mode_raw) {
                print_raw_meta(files[i], &one);
            } else {
                if (file_len > 1) {
                    printf("== %s ==\n", files[i]);
                }
                char *s = tobi_da_emit(data[i], lens[i], &one.meta);
                fputs(s, stdout);
                free(s);
            }
            print_diags(&one.diag, plain);
            if ((strict && tobi_diag_has_error(&one.diag)) || !ok) {
                fail_any = 1;
            }
            tobi_parse_result_free(&one);
        }
        for (size_t i = 0; i < file_len; i++) {
            free(data[i]);
        }
        free(lens);
        free(data);
        free(files);
        return fail_any ? 1 : 0;
    }
    tobi_parse_result res;
    int ok = 0;
    if (file_len == 1) {
        ok = tobi_parse(data[0], lens[0], strict, &res);
        apply_source(&res, files[0], 0);
    } else {
        tobi_parse_input *inputs = tobi_xcalloc(file_len, sizeof(inputs[0]));
        for (size_t i = 0; i < file_len; i++) {
            inputs[i].data = data[i];
            inputs[i].len = lens[i];
            inputs[i].name = files[i];
        }
        ok = tobi_parse_multi(inputs, file_len, strict, &res);
        free(inputs);
    }
    if (!res.root) {
        print_diags(&res.diag, plain);
        tobi_parse_result_free(&res);
        for (size_t i = 0; i < file_len; i++) {
            free(data[i]);
        }
        free(lens);
        free(data);
        free(files);
        return 1;
    }
    (void)tobi_cf_recover(res.root, &res.diag);
    if (mode_json) {
        char *s = tobi_js_emit(&res);
        fputs(s, stdout);
        free(s);
    } else if (mode_dot) {
        char *s = tobi_cf_dot(res.root);
        fputs(s, stdout);
        free(s);
    } else {
        char *s = tobi_dc_emit(res.root, &res.diag);
        fputs(s, stdout);
        free(s);
    }
    print_diags(&res.diag, plain);
    int fail = (strict && tobi_diag_has_error(&res.diag)) || !ok;
    tobi_parse_result_free(&res);
    for (size_t i = 0; i < file_len; i++) {
        free(data[i]);
    }
    free(lens);
    free(data);
    free(files);
    return fail ? 1 : 0;
}