/*
 * This file is part of the MicroPython project, http://micropython.org/
 *
 * The MIT License (MIT)
 *
 * Copyright (c) 2013-2016 Damien P. George
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

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

#include "py/builtin.h"
#include "py/compile.h"
#include "py/persistentcode.h"
#include "py/runtime.h"
#include "py/gc.h"
#include "py/parsenumbase.h"
#include "genhdr/mpversion.h"
#ifdef _WIN32
#include "ports/windows/fmode.h"
#endif

#if MICROPY_EMIT_NATIVE && MICROPY_EMIT_RV32
#include "py/asmrv32.h"

static asm_rv32_backend_options_t rv32_options = { 0 };
#endif

// Command line options, with their defaults
static uint emit_opt = MP_EMIT_OPT_NONE;

#if MICROPY_ENABLE_SOURCE_LINE
static bool include_source_lines = true;
#endif

// Heap size of GC heap (if enabled)
// Make it larger on a 64 bit machine, because pointers are larger.
long heap_size = 1024 * 1024 * (sizeof(mp_uint_t) / 4);

static void stdout_print_strn(void *env, const char *str, size_t len) {
    (void)env;
    ssize_t dummy = write(STDOUT_FILENO, str, len);
    (void)dummy;
}

const mp_print_t mp_stdout_print = {NULL, stdout_print_strn};

static void stderr_print_strn(void *env, const char *str, size_t len) {
    (void)env;
    ssize_t dummy = write(STDERR_FILENO, str, len);
    (void)dummy;
}

static const mp_print_t mp_stderr_print = {NULL, stderr_print_strn};

static int compile_and_save(const char *file, const char *output_file, const char *source_file) {
    nlr_buf_t nlr;
    if (nlr_push(&nlr) == 0) {
        mp_lexer_t *lex;
        if (strcmp(file, "-") == 0) {
            lex = mp_lexer_new_from_fd(MP_QSTR__lt_stdin_gt_, STDIN_FILENO, false);
        } else {
            lex = mp_lexer_new_from_file(qstr_from_str(file));
        }

        qstr source_name;
        if (source_file == NULL) {
            source_name = lex->source_name;
        } else {
            source_name = qstr_from_str(source_file);
        }

        #if MICROPY_MODULE___FILE__
        mp_store_global(MP_QSTR___file__, MP_OBJ_NEW_QSTR(source_name));
        #endif

        mp_parse_tree_t parse_tree = mp_parse(lex, MP_PARSE_FILE_INPUT);
        mp_compiled_module_t cm;
        cm.context = m_new_obj(mp_module_context_t);
        cm.arch_flags = 0;
        #if MICROPY_EMIT_NATIVE && MICROPY_EMIT_RV32
        if (mp_dynamic_compiler.native_arch == MP_NATIVE_ARCH_RV32IMC && mp_dynamic_compiler.backend_options != NULL) {
            cm.arch_flags = ((asm_rv32_backend_options_t *)mp_dynamic_compiler.backend_options)->allowed_extensions;
        }
        #endif

        mp_compile_to_raw_code(&parse_tree, source_name, false, &cm);

        if ((output_file != NULL && strcmp(output_file, "-") == 0) ||
            (output_file == NULL && strcmp(file, "-") == 0)) {
            mp_raw_code_save(&cm, (mp_print_t *)&mp_stdout_print);
        } else {
            vstr_t vstr;
            vstr_init(&vstr, 16);
            if (output_file == NULL) {
                vstr_add_str(&vstr, file);
                vstr_cut_tail_bytes(&vstr, 2);
                vstr_add_str(&vstr, "mpy");
            } else {
                vstr_add_str(&vstr, output_file);
            }

            mp_raw_code_save_file(&cm, qstr_from_strn(vstr.buf, vstr.len));
            vstr_clear(&vstr);
        }

        nlr_pop();
        return 0;
    } else {
        // uncaught exception
        mp_obj_print_exception(&mp_stderr_print, (mp_obj_t)nlr.ret_val);
        return 1;
    }
}

static int usage(char **argv) {
    printf(
        "usage: %s [<opts>] [-X <implopt>] [--] <input filename>\n"
        "Options:\n"
        "--version : show version information\n"
        "-o : output file for compiled bytecode (defaults to input filename with .mpy extension, or stdout if input is stdin)\n"
        "-s : source filename to embed in the compiled bytecode (defaults to input file)\n"
        "-v : verbose (trace various operations); can be multiple\n"
        "-O[N] : apply bytecode optimizations of level N\n"
        "\n"
        "Target specific options:\n"
        "-msmall-int-bits=number : set the maximum bits used to encode a small-int\n"
        "-march=<arch> : set architecture for native emitter;\n"
        "                x86, x64, armv6, armv6m, armv7m, armv7em, armv7emsp,\n"
        "                armv7emdp, xtensa, xtensawin, rv32imc, rv64imc, host, debug\n"
        "-march-flags=<flags> : set architecture-specific flags (can be either a dec/hex/bin value or a comma-separated flags string)\n"
        "                       supported flags for rv32imc: zba, zcmp\n"
        "\n"
        "Implementation specific options:\n", argv[0]
        );
    int impl_opts_cnt = 0;
    printf(
        #if MICROPY_EMIT_NATIVE
        "  emit={bytecode,native,viper} -- set the default code emitter\n"
        #else
        "  emit=bytecode -- set the default code emitter\n"
        #endif
        );
    impl_opts_cnt++;
    printf(
        "  heapsize=<n> -- set the heap size for the GC (default %ld)\n"
        , heap_size);
    impl_opts_cnt++;
    #if MICROPY_ENABLE_SOURCE_LINE
    printf(
        "  source-lines    -- include source line numbers (default)\n"
        "  no-source-lines -- exclude source line numbers\n");
    impl_opts_cnt += 2;
    #endif

    if (impl_opts_cnt == 0) {
        printf("  (none)\n");
    }

    return 1;
}

// Process options which set interpreter init options
static void pre_process_options(int argc, char **argv) {
    for (int a = 1; a < argc; a++) {
        if (argv[a][0] == '-') {
            if (strcmp(argv[a], "-X") == 0) {
                if (a + 1 >= argc) {
                    exit(usage(argv));
                }
                if (strcmp(argv[a + 1], "emit=bytecode") == 0) {
                    emit_opt = MP_EMIT_OPT_BYTECODE;
                #if MICROPY_EMIT_NATIVE
                } else if (strcmp(argv[a + 1], "emit=native") == 0) {
                    emit_opt = MP_EMIT_OPT_NATIVE_PYTHON;
                } else if (strcmp(argv[a + 1], "emit=viper") == 0) {
                    emit_opt = MP_EMIT_OPT_VIPER;
                #endif
                #if MICROPY_ENABLE_SOURCE_LINE
                } else if (strcmp(argv[a + 1], "source-lines") == 0) {
                    // Allow excluding source lines for debug builds.
                    include_source_lines = true;
                } else if (strcmp(argv[a + 1], "no-source-lines") == 0) {
                    // Allow excluding source lines for debug builds.
                    include_source_lines = false;
                #endif
                } else if (strncmp(argv[a + 1], "heapsize=", sizeof("heapsize=") - 1) == 0) {
                    char *end;
                    heap_size = strtol(argv[a + 1] + sizeof("heapsize=") - 1, &end, 0);
                    // Don't bring unneeded libc dependencies like tolower()
                    // If there's 'w' immediately after number, adjust it for
                    // target word size. Note that it should be *before* size
                    // suffix like K or M, to avoid confusion with kilowords,
                    // etc. the size is still in bytes, just can be adjusted
                    // for word size (taking 32bit as baseline).
                    bool word_adjust = false;
                    if ((*end | 0x20) == 'w') {
                        word_adjust = true;
                        end++;
                    }
                    if ((*end | 0x20) == 'k') {
                        heap_size *= 1024;
                    } else if ((*end | 0x20) == 'm') {
                        heap_size *= 1024 * 1024;
                    }
                    if (word_adjust) {
                        heap_size = heap_size * MP_BYTES_PER_OBJ_WORD / 4;
                    }
                } else {
                    exit(usage(argv));
                }
                a++;
            }
        }
    }
}

static char *backslash_to_forwardslash(char *path) {
    for (char *p = path; p != NULL && *p != '\0'; ++p) {
        if (*p == '\\') {
            *p = '/';
        }
    }
    return path;
}

// This will need to be reworked in case mpy-cross needs to set more bits than
// what `unsigned long` can fit.
static bool parse_integer(const char *value, unsigned long *integer) {
    assert(value && "Attempting to parse a NULL string");
    assert(integer && "Attempting to store into a NULL integer buffer");

    size_t value_length = strlen(value);
    int base = 0;
    size_t skip = mp_parse_num_base(value, value_length, &base);
    // These can trip strtoul up.
    if (base < 2 || value_length == skip || value[skip] == '+') {
        return false;
    }
    errno = 0;
    char *end = NULL;
    *integer = strtoul(value + skip, &end, base);
    if (end != (value + value_length) || errno != 0) {
        return false;
    }
    return true;
}

#if MICROPY_EMIT_NATIVE && MICROPY_EMIT_RV32
static bool parse_rv32_flags_string(const char *source, unsigned long *flags) {
    assert(source && "Flag arguments string is NULL.");
    assert(flags && "Collected flags pointer is NULL.");

    const char *current = source;
    const char *end = source + strlen(source);
    unsigned long collected_flags = 0;
    while (current < end) {
        const char *separator = strchr(current, ',');
        if (separator == NULL) {
            separator = end;
        }
        ptrdiff_t length = separator - current;
        if (length == (sizeof("zba") - 1) && memcmp(current, "zba", length) == 0) {
            collected_flags |= RV32_EXT_ZBA;
        } else if (length == (sizeof("zcmp") - 1) && memcmp(current, "zcmp", length) == 0) {
            collected_flags |= RV32_EXT_ZCMP;
        } else {
            return false;
        }
        current = separator + 1;
    }
    *flags = collected_flags;
    return collected_flags != 0;
}
#endif

MP_NOINLINE int main_(int argc, char **argv) {
    pre_process_options(argc, argv);

    char *heap = malloc(heap_size);
    gc_init(heap, heap + heap_size);

    mp_init();
    #ifdef _WIN32
    set_fmode_binary();
    #endif

    #if MICROPY_EMIT_NATIVE
    // Set default emitter options
    MP_STATE_VM(default_emit_opt) = emit_opt;
    #else
    (void)emit_opt;
    #endif

    // set default compiler configuration
    mp_dynamic_compiler.small_int_bits = 31;
    // don't support native emitter unless -march is specified
    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_NONE;
    mp_dynamic_compiler.nlr_buf_num_regs = 0;
    mp_dynamic_compiler.backend_options = NULL;
    #if MICROPY_ENABLE_SOURCE_LINE
    mp_dynamic_compiler.include_source_lines = include_source_lines;
    #endif

    const char *input_file = NULL;
    const char *output_file = NULL;
    const char *source_file = NULL;
    bool option_parsing_active = true;
    const char *arch_flags = NULL;

    // parse main options
    for (int a = 1; a < argc; a++) {
        if (option_parsing_active && argv[a][0] == '-' && argv[a][1] != '\0') {
            if (strcmp(argv[a], "-X") == 0) {
                a += 1;
            } else if (strcmp(argv[a], "--version") == 0) {
                printf(MICROPY_BANNER_NAME_AND_VERSION
                    "; mpy-cross emitting mpy v" MP_STRINGIFY(MPY_VERSION) "." MP_STRINGIFY(MPY_SUB_VERSION) "\n");
                return 0;
            } else if (strcmp(argv[a], "-v") == 0) {
                // This verbose option doesn't currently do anything.
            } else if (strncmp(argv[a], "-O", 2) == 0) {
                if (unichar_isdigit(argv[a][2])) {
                    MP_STATE_VM(mp_optimise_value) = argv[a][2] & 0xf;
                } else {
                    MP_STATE_VM(mp_optimise_value) = 0;
                    for (char *p = argv[a] + 1; *p && *p == 'O'; p++, MP_STATE_VM(mp_optimise_value)++) {;
                    }
                }
            } else if (strcmp(argv[a], "-o") == 0) {
                if (a + 1 >= argc) {
                    exit(usage(argv));
                }
                a += 1;
                output_file = argv[a];
            } else if (strcmp(argv[a], "-s") == 0) {
                if (a + 1 >= argc) {
                    exit(usage(argv));
                }
                a += 1;
                source_file = backslash_to_forwardslash(argv[a]);
            } else if (strncmp(argv[a], "-msmall-int-bits=", sizeof("-msmall-int-bits=") - 1) == 0) {
                char *end;
                mp_dynamic_compiler.small_int_bits =
                    strtol(argv[a] + sizeof("-msmall-int-bits=") - 1, &end, 0);
                if (*end) {
                    return usage(argv);
                }
                // TODO check that small_int_bits is within range of host's capabilities
            } else if (strncmp(argv[a], "-march=", sizeof("-march=") - 1) == 0) {
                const char *arch = argv[a] + sizeof("-march=") - 1;
                if (strcmp(arch, "x86") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_X86;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_X86;
                } else if (strcmp(arch, "x64") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_X64;
                    mp_dynamic_compiler.nlr_buf_num_regs = MAX(MICROPY_NLR_NUM_REGS_X64, MICROPY_NLR_NUM_REGS_X64_WIN);
                } else if (strcmp(arch, "armv6") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV6;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP;
                } else if (strcmp(arch, "armv6m") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV6M;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP; // need to be conservative so this code can run on armv7emdp
                } else if (strcmp(arch, "armv7m") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV7M;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP;
                } else if (strcmp(arch, "armv7em") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV7EM;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP;
                } else if (strcmp(arch, "armv7emsp") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV7EMSP;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP;
                } else if (strcmp(arch, "armv7emdp") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV7EMDP;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP;
                } else if (strcmp(arch, "xtensa") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_XTENSA;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_XTENSA;
                } else if (strcmp(arch, "xtensawin") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_XTENSAWIN;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_XTENSAWIN;
                } else if (strcmp(arch, "rv32imc") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_RV32IMC;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_RV32I;
                } else if (strcmp(arch, "rv64imc") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_RV64IMC;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_RV64I;
                } else if (strcmp(arch, "debug") == 0) {
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_DEBUG;
                    mp_dynamic_compiler.nlr_buf_num_regs = 0;
                } else if (strcmp(arch, "host") == 0) {
                    #if defined(__i386__) || defined(_M_IX86)
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_X86;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_X86;
                    #elif defined(__x86_64__) || defined(_M_X64)
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_X64;
                    mp_dynamic_compiler.nlr_buf_num_regs = MAX(MICROPY_NLR_NUM_REGS_X64, MICROPY_NLR_NUM_REGS_X64_WIN);
                    #elif defined(__arm__) && !defined(__thumb2__)
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_ARMV6;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_ARM_THUMB_FP;
                    #elif defined(__riscv) && (__riscv_xlen == 64)
                    mp_dynamic_compiler.native_arch = MP_NATIVE_ARCH_RV64IMC;
                    mp_dynamic_compiler.nlr_buf_num_regs = MICROPY_NLR_NUM_REGS_RV64I;
                    #else
                    mp_printf(&mp_stderr_print, "unable to determine host architecture for -march=host\n");
                    exit(1);
                    #endif
                } else {
                    return usage(argv);
                }
            } else if (strncmp(argv[a], "-march-flags=", sizeof("-march-flags=") - 1) == 0) {
                arch_flags = argv[a] + sizeof("-march-flags=") - 1;
            } else if (strcmp(argv[a], "--") == 0) {
                option_parsing_active = false;
            } else {
                return usage(argv);
            }
        } else {
            if (input_file != NULL) {
                mp_printf(&mp_stderr_print, "multiple input files\n");
                exit(1);
            }
            input_file = backslash_to_forwardslash(argv[a]);
        }
    }

    if (arch_flags && mp_dynamic_compiler.native_arch != MP_NATIVE_ARCH_NONE) {
        bool processed = false;
        #if MICROPY_EMIT_NATIVE && MICROPY_EMIT_RV32
        if (mp_dynamic_compiler.native_arch == MP_NATIVE_ARCH_RV32IMC) {
            mp_dynamic_compiler.backend_options = (void *)&rv32_options;
            unsigned long raw_flags = 0;
            if (parse_integer(arch_flags, &raw_flags) || parse_rv32_flags_string(arch_flags, &raw_flags)) {
                if ((raw_flags & ~((unsigned long)RV32_EXT_ALL)) == 0) {
                    rv32_options.allowed_extensions = (uint8_t)raw_flags;
                    processed = true;
                }
            }
        }
        #endif
        if (!processed) {
            mp_printf(&mp_stderr_print, "unrecognised arch flags\n");
            exit(1);
        }
    }

    #if MICROPY_EMIT_NATIVE
    if ((MP_STATE_VM(default_emit_opt) == MP_EMIT_OPT_NATIVE_PYTHON
         || MP_STATE_VM(default_emit_opt) == MP_EMIT_OPT_VIPER)
        && mp_dynamic_compiler.native_arch == MP_NATIVE_ARCH_NONE) {
        mp_printf(&mp_stderr_print, "arch not specified\n");
        exit(1);
    }
    #endif

    if (input_file == NULL) {
        mp_printf(&mp_stderr_print, "no input file\n");
        exit(1);
    }

    int ret = compile_and_save(input_file, output_file, source_file);

    mp_deinit();

    return ret & 0xff;
}

int main(int argc, char **argv) {
    mp_cstack_init_with_sp_here(40000 * (sizeof(void *) / 4));
    return main_(argc, argv);
}

mp_import_stat_t mp_import_stat(const char *path) {
    (void)path;
    return MP_IMPORT_STAT_NO_EXIST;
}

void nlr_jump_fail(void *val) {
    fprintf(stderr, "FATAL: uncaught NLR %p\n", val);
    exit(1);
}
