#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>

#ifndef TOBI_VERSION
#define TOBI_VERSION "0.0.0"
#endif

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"
          "  --graph K   print DOT graph: cfg, ast, namespace\n"
          "  --graph-labels K\n"
          "              set graph labels: short, full, none\n"
          "  --raw       print input/table metadata and aml offsets\n"
          "  -o FILE     write primary output to FILE instead of stdout\n"
          "  --strict    treat malformed or unsupported constructs as hard errors\n"
          "  --plain     disable coloured diagnostics\n"
          "  --help      print usage\n"
          "  --version   print version\n"
          "\n"
          "Use '-' as an input path to read AML from stdin.\n", f);
}

static int append_file(const char ***files, size_t *file_len, size_t *file_cap, const char *path) {
    if (*file_len == *file_cap) {
        *file_cap = tobi_xgrow_cap(*file_cap, *file_len + 1u, 4u);
        *files = tobi_xrealloc((void *)*files, tobi_xmul_size(*file_cap, sizeof((*files)[0])));
    }
    (*files)[(*file_len)++] = path;
    return 1;
}

static const char *source_label(const char *path) {
    return path && strcmp(path, "-") == 0 ? "<stdin>" : path;
}

static int read_stream(FILE *fp, const char *label, int close_stream, uint8_t **data, size_t *len) {
    size_t cap = 4096;
    size_t used = 0;
    uint8_t *buf = tobi_xmalloc(cap);
    for (;;) {
        if (used == cap) {
            cap = tobi_xgrow_cap(cap, tobi_xadd_size(used, 1u), 4096u);
            buf = tobi_xrealloc(buf, cap);
        }
        size_t room = cap - used;
        size_t got = fread(buf + used, 1, room, fp);
        if (got > room) {
            fprintf(stderr, "tobi: %s: read exceeded buffer\n", label);
            free(buf);
            if (close_stream) {
                fclose(fp);
            }
            return 0;
        }
        used += got;
        if (got < room) {
            if (ferror(fp)) {
                fprintf(stderr, "tobi: %s: read failed: %s\n", label, strerror(errno));
                free(buf);
                if (close_stream) {
                    fclose(fp);
                }
                return 0;
            }
            break;
        }
    }
    if (close_stream && fclose(fp) != 0) {
        fprintf(stderr, "tobi: %s: close failed: %s\n", label, strerror(errno));
        free(buf);
        return 0;
    }
    *data = buf;
    *len = used;
    return 1;
}

static int read_file(const char *path, uint8_t **data, size_t *len) {
    if (strcmp(path, "-") == 0) {
        return read_stream(stdin, "<stdin>", 0, data, 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 FILE *open_output(const char *path) {
    if (!path || strcmp(path, "-") == 0) {
        return stdout;
    }
    FILE *fp = fopen(path, "wb");
    if (!fp) {
        fprintf(stderr, "tobi: %s: %s\n", path, strerror(errno));
        return NULL;
    }
    return fp;
}

static int close_output(FILE *fp, const char *path) {
    if (!fp || fp == stdout) {
        return 1;
    }
    if (fclose(fp) != 0) {
        fprintf(stderr, "tobi: %s: close failed: %s\n", path, strerror(errno));
        return 0;
    }
    return 1;
}

static int parse_graph_kind(const char *s, tobi_graph_kind *kind) {
    if (strcmp(s, "cfg") == 0) {
        *kind = TOBI_GRAPH_CFG;
        return 1;
    }
    if (strcmp(s, "ast") == 0) {
        *kind = TOBI_GRAPH_AST;
        return 1;
    }
    if (strcmp(s, "namespace") == 0 || strcmp(s, "ns") == 0) {
        *kind = TOBI_GRAPH_NAMESPACE;
        return 1;
    }
    return 0;
}

static int parse_graph_labels(const char *s, tobi_graph_labels *labels) {
    if (strcmp(s, "short") == 0) {
        *labels = TOBI_GRAPH_LABEL_SHORT;
        return 1;
    }
    if (strcmp(s, "full") == 0) {
        *labels = TOBI_GRAPH_LABEL_FULL;
        return 1;
    }
    if (strcmp(s, "none") == 0) {
        *labels = TOBI_GRAPH_LABEL_NONE;
        return 1;
    }
    return 0;
}

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;
    }
    free(meta->source);
    meta->source = tobi_xstrdup(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(FILE *out, const char *input, const tobi_parse_result *res) {
    fprintf(out, "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.source : "raw", 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);
}

static void free_inputs(uint8_t **data, size_t *lens, size_t file_len, const char **files) {
    (void)lens;
    for (size_t i = 0; i < file_len; i++) {
        free(data[i]);
    }
    free(lens);
    free(data);
    free(files);
}

int main(int argc, char **argv) {
    int mode_dis = 0;
    int mode_json = 0;
    int mode_graph = 0;
    int mode_raw = 0;
    int strict = 0;
    int plain = 0;
    int stdin_count = 0;
    const char *output_path = NULL;
    tobi_graph_kind graph_kind = TOBI_GRAPH_CFG;
    tobi_graph_labels graph_labels = TOBI_GRAPH_LABEL_SHORT;
    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_graph = 1;
            graph_kind = TOBI_GRAPH_CFG;
        } else if (strcmp(argv[i], "--graph") == 0) {
            if (i + 1 >= argc) {
                fprintf(stderr, "tobi: --graph requires cfg, ast, or namespace\n");
                free(files);
                return 2;
            }
            if (!parse_graph_kind(argv[++i], &graph_kind)) {
                fprintf(stderr, "tobi: unknown graph kind %s\n", argv[i]);
                free(files);
                return 2;
            }
            mode_graph = 1;
        } else if (strcmp(argv[i], "--graph-labels") == 0) {
            if (i + 1 >= argc) {
                fprintf(stderr, "tobi: --graph-labels requires short, full, or none\n");
                free(files);
                return 2;
            }
            if (!parse_graph_labels(argv[++i], &graph_labels)) {
                fprintf(stderr, "tobi: unknown graph label mode %s\n", argv[i]);
                free(files);
                return 2;
            }
        } else if (strcmp(argv[i], "--raw") == 0) {
            mode_raw = 1;
        } else if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--output") == 0) {
            if (i + 1 >= argc) {
                fprintf(stderr, "tobi: %s requires a file\n", argv[i]);
                free(files);
                return 2;
            }
            output_path = argv[++i];
        } else if (strcmp(argv[i], "--strict") == 0) {
            strict = 1;
        } else if (strcmp(argv[i], "--plain") == 0) {
            plain = 1;
        } else if (strcmp(argv[i], "-") == 0) {
            stdin_count++;
            (void)append_file(&files, &file_len, &file_cap, argv[i]);
        } else if (argv[i][0] == '-') {
            fprintf(stderr, "tobi: unknown option %s\n", argv[i]);
            usage(stderr);
            return 2;
        } else {
            (void)append_file(&files, &file_len, &file_cap, argv[i]);
        }
    }
    if (file_len == 0) {
        usage(stderr);
        free(files);
        return 2;
    }
    if (stdin_count > 1) {
        fprintf(stderr, "tobi: stdin can only be used once\n");
        free(files);
        return 2;
    }
    if ((mode_dis + mode_json + mode_graph + mode_raw) > 1) {
        fprintf(stderr, "tobi: choose only one output mode\n");
        free(files);
        return 2;
    }
    FILE *outfp = open_output(output_path);
    if (!outfp) {
        free(files);
        return 1;
    }
    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])) {
            (void)close_output(outfp, output_path);
            free_inputs(data, lens, file_len, files);
            return 1;
        }
    }
    if (mode_dis || mode_raw) {
        int fail_any = 0;
        for (size_t i = 0; i < file_len; i++) {
            const char *label = source_label(files[i]);
            tobi_parse_result one;
            int ok = tobi_parse(data[i], lens[i], strict, &one);
            apply_source(&one, label, 0);
            if (mode_raw) {
                print_raw_meta(outfp, label, &one);
            } else {
                if (file_len > 1) {
                    fprintf(outfp, "== %s ==\n", label);
                }
                char *s = tobi_da_emit(data[i], lens[i], &one.meta);
                fputs(s, outfp);
                free(s);
            }
            print_diags(&one.diag, plain);
            if ((strict && tobi_diag_has_error(&one.diag)) || !ok) {
                fail_any = 1;
            }
            tobi_parse_result_free(&one);
        }
        int output_ok = close_output(outfp, output_path);
        free_inputs(data, lens, file_len, files);
        if (!output_ok) {
            return 1;
        }
        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, source_label(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 = source_label(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);
        (void)close_output(outfp, output_path);
        free_inputs(data, lens, file_len, files);
        return 1;
    }
    (void)tobi_cf_recover(res.root, &res.diag);
    if (mode_json) {
        char *s = tobi_js_emit(&res);
        fputs(s, outfp);
        free(s);
    } else if (mode_graph) {
        char *s = tobi_graph_dot_labels(res.root, &res.ns, graph_kind, graph_labels);
        fputs(s, outfp);
        free(s);
    } else {
        char *s = tobi_dc_emit(res.root, &res.diag);
        fputs(s, outfp);
        free(s);
    }
    print_diags(&res.diag, plain);
    int fail = (strict && tobi_diag_has_error(&res.diag)) || !ok;
    tobi_parse_result_free(&res);
    int output_ok = close_output(outfp, output_path);
    free_inputs(data, lens, file_len, files);
    if (!output_ok) {
        return 1;
    }
    return fail ? 1 : 0;
}