#include <errno.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>

#include "apfl.h"
#include "bytecode.h"
#include "compile.h"
#include "format.h"
#include "gc.h"
#include "strings.h"

struct c_writer_data {
    struct apfl_io_writer w;
    size_t written_len;
};

static const unsigned char hex[] = {
    '0', '1', '2', '3', '4', '5', '6', '7',
    '8', '9', 'A', 'B', 'C', 'D', 'E', 'F',
};

static bool
write_as_c_source(void *opaque, const unsigned char *buf, size_t len)
{
    struct c_writer_data *data = opaque;
    for (size_t i = 0; i < len; i++) {
        if (data->written_len % 16 == 0) {
            if (!apfl_io_write_byte(data->w, '\n')) {
                return false;
            }
        }

        unsigned char out[] = {'0', 'x', hex[(buf[i] >> 4) & 0x0F], hex[buf[i] & 0x0F], ','};

        if (!apfl_io_write_string_view(
            data->w,
            (struct apfl_string_view) { .bytes = out, .len = sizeof(out), }
        )) {
            return false;
        }

        (data->written_len)++;
    }

    return true;
}

static bool
c_source_prelude(struct apfl_io_writer w)
{
    return apfl_io_write_string(
        w,
        "#include <stddef.h>\n"
        // TODO: This is a bit hackish, the generated C code should include
        // apfl.h instead of declaring the string view struct itself. However,
        // since the code is generated, it does not live in the same directory
        // as apfl.h, so our usual `#include "apfl.h"` doesn't work.
        "struct apfl_string_view {\n"
        "  const unsigned char *bytes;\n"
        "  size_t len;\n"
        "};\n"
        "static const unsigned char bytecode[] = {"
    );
}

static bool
c_source_epilogue(struct apfl_io_writer w, const char *function_name)
{
    return apfl_io_write_string(w,
        "\n};\n"
        "struct apfl_string_view\n"
    )
        && apfl_io_write_string(w, function_name)
        && apfl_io_write_string(w,
            "(void)\n"
            "{\n"
            "    return (struct apfl_string_view) {\n"
            "        .bytes = bytecode,\n"
            "        .len = sizeof(bytecode),\n"
            "    };\n"
            "}\n"
        );
}

struct roots {
    struct instruction_list *ilist;
    struct apfl_string *name;
};

static void
getroots(void *own_opaque, gc_visitor visitor, void *visitor_opaque)
{
    struct roots *roots = own_opaque;
    if (roots->ilist != NULL) {
        visitor(visitor_opaque, GC_OBJECT_FROM(roots->ilist, GC_TYPE_INSTRUCTIONS));
    }
    if (roots->name != NULL) {
        visitor(visitor_opaque, GC_OBJECT_FROM(roots->name, GC_TYPE_STRING));
    }
}

static const char *
get_next_string_from_arg(const char *progname, const char ***argpp)
{
    if (**argpp == NULL) {
        fprintf(stderr, "Usage: %s [-c function_name] input output\n", progname);
        return NULL;
    }

    const char *filename = **argpp;
    (*argpp)++;
    return filename;
}

static FILE *
openfile(const char *progname, const char *filename, const char *mode)
{
    FILE *f = fopen(filename, mode);
    if (f == NULL) {
        fprintf(stderr, "%s: Could not open %s: %s\n", progname, filename, strerror(errno));
        return NULL;
    }
    return f;
}

static void
closefile(FILE *f)
{
    if (f != NULL) {
        fclose(f);
    }
}

static bool
compile(
    const char *progname,
    struct gc *gc,
    struct apfl_io_reader r,
    struct instruction_list *ilist
) {
    apfl_tokenizer_ptr tokenizer = apfl_tokenizer_new(gc->allocator, apfl_io_reader_as_source_reader(&r));
    if (tokenizer == NULL) {
        return false;
    }

    apfl_parser_ptr parser = apfl_parser_new(
        gc->allocator,
        apfl_tokenizer_as_token_source(tokenizer)
    );
    if (parser == NULL) {
        apfl_tokenizer_destroy(tokenizer);
        fprintf(stderr, "%s: Failed to init parser\n", progname);
        return false;
    }

    struct apfl_error error;
    if (!apfl_compile_whole_file(
        gc,
        parser,
        &error,
        ilist
    )) {
        fprintf(stderr, "%s: Could not compile file:\n", progname);
        apfl_error_print(error, stderr);

        apfl_parser_destroy(parser);
        apfl_tokenizer_destroy(tokenizer);
        return false;
    }

    apfl_parser_destroy(parser);
    apfl_tokenizer_destroy(tokenizer);
    return true;
}

int
main(int argc, const char *argv[])
{
    (void)argc;
    const char **argp = argv+1;

    FILE *in = NULL;
    FILE *out = NULL;

    const char *c_function_name = NULL;
    if (*argp != NULL && apfl_string_eq(*argp, "-c")) {
        argp++;
        if ((c_function_name = get_next_string_from_arg(argv[0], &argp)) == NULL) {
            goto error_before_gc;
        }
    }

    const char *name_in = NULL;
    const char *name_out = NULL;
    if (
        (name_in = get_next_string_from_arg(argv[0], &argp)) == NULL
        || (name_out = get_next_string_from_arg(argv[0], &argp)) == NULL
    ) {
        goto error_before_gc;
    }

    int rv = 1;

    if (
        (in = openfile(argv[0], name_in, "rb")) == NULL
        || (out = openfile(argv[0], name_out, "wb")) == NULL
    ) {
        goto error_before_gc;
    }

    struct apfl_io_reader r = apfl_io_file_reader(in);
    struct apfl_io_writer w_raw = apfl_io_file_writer(out);

    struct roots roots = {
        .ilist = NULL,
        .name = NULL,
    };

    struct gc gc;
    apfl_gc_init(&gc, apfl_allocator_default(), getroots, &roots);

    if ((roots.name = apfl_string_copy_to_new_gc_string(
        &gc,
        apfl_string_view_from(name_in)
    )) == NULL) {
        goto error;
    }

    if ((roots.ilist = apfl_instructions_new(&gc, 1, roots.name)) == NULL) {
        goto error;
    }

    if (!compile(argv[0], &gc, r, roots.ilist)) {
        goto error;
    }

    if (c_function_name != NULL && !c_source_prelude(w_raw)) {
        fprintf(stderr, "%s: IO error\n", argv[0]);
        goto error;
    }

    struct apfl_io_writer bytecode_writer = w_raw;
    struct c_writer_data c_writer_data = {
        .w = w_raw,
        .written_len = 0,
    };
    if (c_function_name != NULL) {
        bytecode_writer = (struct apfl_io_writer) {
            .opaque = &c_writer_data,
            .write = write_as_c_source,
        };
    }

    if (!apfl_bytecode_serialize(
        gc.allocator,
        bytecode_writer,
        roots.ilist
    )) {
        fprintf(stderr, "%s: Could not serialize bytecode\n", argv[0]);
        goto error;
    }

    if (c_function_name != NULL && !c_source_epilogue(w_raw, c_function_name)) {
        fprintf(stderr, "%s: IO error\n", argv[0]);
        goto error;
    }

    rv = 0;
error:
    apfl_gc_deinit(&gc);
error_before_gc:
    closefile(in);
    closefile(out);
    return rv;
}