#!/usr/bin/env python3

import argparse
import glob
import os
import re
import signal
import subprocess
import sys
import shlex

from binascii import hexlify
from ctypes.util import find_library
from functools import reduce
from itertools import chain
try:
    from shutil import which
except ImportError:
    from backports.shutil_which import which

from elftools.elf.dynamic import DynamicSegment
from elftools.elf.elffile import ELFFile

search_prefixes = []
real_argv0 = os.environ.get("REAL_ARGV0", sys.argv[0])


def shlex_join(split_command):
    return ' '.join(shlex.quote(arg) for arg in split_command)


def log_error(msg):
    sys.stderr.write(msg + "\n")


def relative(path):
    return os.path.relpath(path, ".")


def run(command, override={}):
    if is_executable(command[0]):
        command = wrap(command)

    global log_commands
    if log_commands:
        cwd = os.getcwd().rstrip("/")
        if command[0].startswith(cwd):
            program_path = "." + command[0][len(cwd):]
        else:
            program_path = command[0]
        sys.stderr.write("{}\n\n".format(
            " \\\n  ".join([program_path] + command[1:])))

    signal.signal(signal.SIGINT, signal.SIG_IGN)
    environment = dict(os.environ)
    environment.update(override)
    p = subprocess.Popen(command,
                         preexec_fn=lambda: signal.signal(signal.SIGINT,
                                                          signal.SIG_DFL),
                         env=environment)
    if p.wait() != 0:
        log_error("The following command exited with {}:\n{}".format(
            p.returncode, shlex_join(command)))
        sys.exit(p.returncode)


def is_executable(path):
    with open(path, "rb") as program:
        return program.read(4) == b"\x7fELF"


def is_dynamic(path):
    with open(path, "rb") as input_file:
        return len([segment
                    for segment
                    in ELFFile(input_file).iter_segments()
                    if segment.header.p_type == "PT_DYNAMIC"]) != 0


def interleave(base, repeat):
    return list(sum(zip([repeat] * len(base), base), ()))

# Use in case different version of pyelftools might give str or bytes


def to_string(obj):
    if type(obj) is str:
        return obj
    elif type(obj) is bytes:
        return obj.decode("utf-8")


def get_elf_needed(path):
    with open(path, "rb") as elf_file:
        segments = [segment
                    for segment
                    in ELFFile(elf_file).iter_segments()
                    if type(segment) is DynamicSegment]

        if len(segments) == 1:
            needed = [to_string(tag.needed)
                      for tag
                      in segments[0].iter_tags()
                      if tag.entry.d_tag == "DT_NEEDED"]

            runpath = [tag.runpath
                       for tag
                       in segments[0].iter_tags()
                       if tag.entry.d_tag == "DT_RUNPATH"]

            assert len(runpath) < 2

            if not runpath:
                return needed
            else:
                runpaths = [path.replace("$ORIGIN",
                                         os.path.dirname(path))
                            for path
                            in runpath[0].split(":")]
                absolute_needed = []
                for lib in needed:
                    found = False
                    for runpath in runpaths:
                        full_path = os.path.join(runpath, lib)
                        if os.path.isfile(full_path):
                            absolute_needed.append(full_path)
                            found = True
                            break

                    if not found:
                        absolute_needed.append(lib)

                return absolute_needed

        else:
            return []


def get_command(command):
    global search_path
    path = which(command, path=search_path)
    if not path:
        log_error("""Couldn't find "{}" in "{}".""".format(
            command, search_path))
        assert False
    return os.path.abspath(path)


def build_opt_args(args):
    analysis_libraries = []
    all = set()
    for prefix in search_prefixes:
        analyses_path = os.path.join(prefix, "lib", "revng", "analyses")
        if not os.path.isdir(analyses_path):
            continue

        # Enumerate all the libraries containing analyses
        for library in glob.glob(os.path.join(analyses_path, "*.so")):
            basename = os.path.basename(library)
            if basename not in all:
                analysis_libraries.append(library)
                all.add(basename)

    # Identify all the libraries that are dependencies of other libraries, i.e.,
    # non-roots in the dependencies tree. Note that circular dependencies are not
    # allowed.
    provided = set(chain.from_iterable([get_elf_needed(path)
                                        for path
                                        in analysis_libraries]))

    # Path to libraries representing the roots in the dependency tree
    roots = [relative(os.path.join(analyses_path, name))
             for name
             in all - provided]

    prefix = []
    libasan = [name
               for name
               in provided
               if ("libasan." in name
                   or "libclang_rt.asan" in name)]
    if libasan:
        original_asan_options = os.environ.get("ASAN_OPTIONS", "")
        if original_asan_options:
            asan_options = dict([option.split("=")
                                 for option
                                 in original_asan_options.split(":")])
        else:
            asan_options = dict()
        asan_options["abort_on_error"] = "1"
        new_asan_options = ":".join(["=".join(option)
                                     for option
                                     in asan_options.items()])

        # Use `sh` instead of `env` since `env` sometimes is not a real executable
        # but a shebang script spawning /usr/bin/coreutils, which makes gdb unhappy
        prefix = [get_command("sh"),
                  "-c",
                  'LD_PRELOAD={} ASAN_OPTIONS={} '
                  'exec "$0" "$@"'.format(libasan[0], new_asan_options)]

    return (prefix + [relative(get_command("opt"))]
            + interleave(roots, "-load")
            + args
            + ["-serialize-model"])


def split_dash_dash(args):
    if not args:
        return [], []

    extra_args = []
    while (len(args) != 0) and (args[0] != "--"):
        extra_args.append(args[0])
        del args[0]

    # Drop the delimiter
    if len(args) != 0:
        del args[0]

    return list(extra_args), args


def get_architecture(path):
    with open(path, "rb") as file:
        file.seek(5)
        arch_id = hexlify(file.read(1))
        file.seek(18)
        arch_id += hexlify(file.read(2))
    arch_id = arch_id.decode("utf-8")

    if arch_id == "012800":
        return "arm"
    elif arch_id == "01b700":
        return "aarch64"
    elif arch_id == "010800":
        return "mipsel"
    elif arch_id == "020008":
        return "mips"
    elif arch_id == "013e00":
        return "x86_64"
    elif arch_id == "010300":
        return "i386"
    elif arch_id == "020016":
        return "s390x"
    elif arch_id == "01b700":
        return "aarch64"
    else:
        log_error("Unknown architecture ID: {}".format(arch_id))
        assert False


def find_file(name, paths):
    for path in paths:
        path = os.path.join(path, name)
        if os.path.isfile(path):
            return path

    log_error("Can't find the following file: {}".format(name))
    assert False


def get_stderr(args):
    with subprocess.Popen(args, stderr=subprocess.PIPE) as process:
        return process.stderr.read()


def build_linking_options(li_csv_path, need_csv_path):
    result = []

    # Force maximum page size
    page_size = 4096
    result.append("-Wl,-z,max-page-size={}".format(page_size))

    # Produce semgment loading options
    min = 0
    max = 0
    with open(li_csv_path, "r") as li_csv_file:
        # Consume CSV header
        li_csv_file.readline()

        for segment in li_csv_file:
            name, start, end = segment.split(",")
            start, end = int(start, 16), int(end, 16)

            # Record max and min
            if end > max:
                max = end

            if (min == 0) or start < min:
                min = start

            # Force section address
            result.append("-Wl,--section-start={}={}".format(name, hex(start)))

    # Force ld.bfd
    result.append("-fuse-ld=bfd")

    # Force a page before the lowest original address for the ELF header
    result.append(
        "-Wl,--section-start=.elfheaderhelper={}".format(hex(min - 1)))

    # Force text to start on the page after all the original program segments
    text_address = page_size * int((max + page_size - 1) / page_size)
    result.append("-Wl,-Ttext-segment={}".format(hex(text_address)))

    # Link required dynamic libraries
    result.append("-Wl,--no-as-needed")

    with open(need_csv_path, "r") as need_csv_file:
        for library in need_csv_file:
            library = library.strip()
            if "/" in library:
                library = library[library.rindex("/") + 1:]

            match = re.match("^lib(.*).so(\.[0-9]+)*$", library)
            if match:
                library = "-l" + match.group(1)
            else:
                library = "-l:" + library

            result.append(library)

    result.append("-Wl,--as-needed")

    return result


def wrap(args):
    global command_prefix
    return command_prefix + args


def register_translate(subparsers):
    parser = subparsers.add_parser("translate",
                                   help="the rev.ng translator",
                                   description="The rev.ng translator.",
                                   prog=real_argv0 + " translate")
    parser.add_argument("-O0", action="store_true", help="Do no optimize.")
    parser.add_argument("-O1", action="store_true", help="Use llc -O2.")
    parser.add_argument("-O2",
                        action="store_true",
                        help="Use llc -O2 and opt -O2.")
    parser.add_argument("--trace",
                        action="store_true",
                        help="Use the tracing version of support.ll.")
    parser.add_argument("-s",
                        "--skip",
                        action="store_true",
                        help="Do not invoke revng-lift.")
    parser.add_argument("-i",
                        "--isolate",
                        action="store_true",
                        help="Enable function isolation.")
    parser.add_argument("--base", help="Load address to employ in lifting.")
    parser.add_argument("-o", "--output", metavar="OUTPUT",
                        help="Output path.")
    parser.add_argument("input", metavar="INPUT", help="The input binary.")


def run_translate(args, post_dash_dash):
    # Ensure there's an input file
    if not args.input:
        log_error("Input file not provided")
        return -1

    # Register optimization level
    optimization_level = 0
    if args.O1:
        optimization_level = 1
    elif args.O2:
        optimization_level = 2

    input = args.input
    executable = args.output if args.output else "{}.translated".format(input)
    output = "{}.ll".format(executable)
    need_csv_path = "{}.need.csv".format(output)
    li_csv_path = "{}.li.csv".format(output)

    # Find a compiler (used for linking)
    compiler = get_command(os.environ.get("CXX", "c++"))

    # Extracted the architecture name
    architecture = get_architecture(input)

    if is_dynamic(input) and architecture != "x86_64":
        log_error("""
Currently, rev.ng can translate dynamically linked binaries only for x86-64.
The provided program is dynamically linked but targets {}.
Try with an input program that is compiled for x86-64, or statically linked.
If you're just interested in lifting to LLVM IR, simply use `revng lift`.
    """.format(architecture).strip())
        return -1

    # Check if tracing is enabled
    if args.trace:
        config = "trace"
    else:
        config = "normal"

    # Build the name of the support.ll file
    support_name = "support-{}-{}.ll".format(architecture, config)

    # Find support.ll
    global script_path
    support_path = find_file(support_name,
                             [os.path.join(script_path,
                                           "..",
                                           "share",
                                           "revng"),
                              script_path])

    # Perform lifting
    if not args.skip:

        lift_options = post_dash_dash or []

        if args.base:
            lift_options += ["--base", args.base]

        run([get_command("revng-lift"),
             "-g", "ll"]
            + lift_options
            + [relative(input), relative(output)])

    # Perform function isolation
    if args.isolate:
        isolated = "{}.isolated".format(executable)
        opt_invocation = build_opt_args(["-S",
                                         "-detect-abi",
                                         "-isolate",
                                         "-invoke-isolated-functions",
                                         relative(output),
                                         "-o", relative(isolated)])
        run(opt_invocation)
        output = isolated

    # Link with support
    linked = "{}.linked.ll".format(output)
    run([get_command("llvm-link"),
         "-S",
         relative(output),
         relative(support_path),
         "-o", relative(linked)])
    output = linked

    # Compile
    object_file = "{}.o".format(output)

    common_llc_options = ["-disable-machine-licm", "-filetype=obj"]
    llc = get_command("llc")
    if optimization_level == 0:
        run([llc,
             "-O0",
             relative(output),
             "-o", relative(object_file)]
            + common_llc_options)
    elif optimization_level == 1:
        run([llc,
             "-O2",
             relative(output),
             "-o", relative(object_file)]
            + common_llc_options)
    elif optimization_level == 2:
        optimized = "{}.opt.ll".format(output)
        run([get_command("opt"),
             "-O2",
             "-S",
             "-enable-pre=false",
             "-enable-load-pre=false",
             relative(output),
             "-o", relative(optimized)])
        run([llc,
             "-O2",
             relative(optimized),
             "-o", relative(object_file)]
            + common_llc_options)

    # Parse .li.csv and .need.csv files
    linking_options = build_linking_options(li_csv_path, need_csv_path)

    no_pie = ["-fno-pie"]

    # Detect -no-pie support
    if b"unrecognized command line" not in get_stderr([compiler, "-no-pie"]):
        no_pie.append("-no-pie")

    run([compiler,
         object_file,
         "-lz", "-lm", "-lrt", "-lpthread",
         "-L", "./",
         "-o", executable]
        + no_pie
        + linking_options,
        {"HARD_FLAGS_IGNORE": "1"})

    # Invoke revng-merge-dynamic
    unpatched_output = "{}.tmp".format(executable)
    os.rename(executable, unpatched_output)
    base_args = ["--base", args.base] if args.base else []
    run([get_command("revng"),
         "merge-dynamic",
         unpatched_output,
         input,
         executable]
        + base_args)
    os.unlink(unpatched_output)

    return 0


def run_cc(args):
    # Collect translate options
    translated_args = []
    if "--" in args:
        translate_args, args = split_dash_dash(args)

    if (not args) or ("--help" in args):
        log_error(f"Usage: {real_argv0} cc [[translate-options] --] compiler "
                  + "[compiler-options]")
        return -1

    res = subprocess.call(args)
    if res != 0:
        return res

    # Are we linking?
    if not ("-c" in args):
        assert "-o" in args

        # Identify the path of the final program
        output = os.path.abspath(args[args.index("-o") + 1])

        original = output + ".original"
        translated = original + ".translated"
        os.rename(output, original)

        result = run_translate(translate_args + [original])
        if result != 0:
            return result

        os.rename(translated, output)

    return res


def main():
    parser = argparse.ArgumentParser(description="The rev.ng driver.")
    parser.add_argument("--version",
                        action="store_true",
                        help="Display version information.")
    parser.add_argument("--verbose",
                        action="store_true",
                        help="Log all executed commands.")
    parser.add_argument("--perf",
                        action="store_true",
                        help="Run programs under perf (for use with hotspot).")
    parser.add_argument("--heaptrack",
                        action="store_true",
                        help="Run programs under heaptrack.")
    parser.add_argument("--gdb",
                        action="store_true",
                        help="Run programs under gdb.")
    parser.add_argument("--lldb",
                        action="store_true",
                        help="Run programs under lldb.")
    parser.add_argument("--valgrind",
                        action="store_true",
                        help="Run programs under valgrind.")
    parser.add_argument("--callgrind",
                        action="store_true",
                        help="Run programs under callgrind.")
    parser.add_argument("--prefix",
                        metavar="PREFIX",
                        help="Additional search prefix.")

    # Get script_path
    global script_path
    script_path = os.path.dirname(os.path.realpath(__file__))

    # Create a custom PATH variable to find programs
    global search_path
    search_path = os.environ["PATH"]
    if search_path:
        search_path = "{}:{}".format(script_path, search_path)
    else:
        search_path = script_path

    subparsers = parser.add_subparsers(dest="command_name",
                                       help='sub-commands help')

    register_translate(subparsers)

    subparsers.add_parser("cc",
                          help="compile, link and translate transparently",
                          add_help=False)
    subparsers.add_parser("opt",
                          help="LLVM's opt with rev.ng passes",
                          add_help=False)

    programs = set()
    prefix = "revng-"
    for path in search_path.split(":"):
        if os.path.isdir(path):
            for program in os.listdir(path):
                if program.startswith(prefix) and which(program, path=search_path):
                    name = program[len(prefix):]
                    subparsers.add_parser(name,
                                          help=f"see {name} --help",
                                          add_help=False)

    global revng_parser
    revng_parser = parser

    # Strip away arguments -- so we can forward them to revng-lift
    base_args, post_dash_dash = split_dash_dash(sys.argv[1:])

    args, unknown_args = parser.parse_known_args(base_args)

    global search_prefixes
    search_prefixes = (args.prefix or []
                       + [os.path.join(script_path, "..")])

    global log_commands
    log_commands = args.verbose

    global command_prefix
    command_prefix = []

    assert (args.gdb + args.lldb + args.valgrind + args.callgrind) <= 1

    if args.gdb:
        command_prefix += ["gdb", "-q", "--args"]

    if args.lldb:
        command_prefix += ["lldb", "--"]

    if args.valgrind:
        command_prefix += ["valgrind"]

    if args.callgrind:
        command_prefix += ["valgrind", "--tool=callgrind"]

    if args.perf:
        command_prefix += ["perf",
                           "record",
                           "--call-graph", "dwarf",
                           "--output=perf.data"]

    if args.heaptrack:
        command_prefix += ["heaptrack"]

    if args.version:
        sys.stdout.write("rev.ng version @VERSION@\n")
        return 0

    return run_command(args, unknown_args, post_dash_dash)


def run_command(args, unknown_args, post_dash_dash):
    command = args.command_name
    if not command:
        log_error("No command specified")
        return 1

    all_args = list(unknown_args)
    if post_dash_dash:
        all_args += ["--"] + post_dash_dash

    # First consider the hardcoded commands
    if command == "opt":
        run(build_opt_args(all_args))
    elif command == "cc":
        return run_cc(all_args)
    elif command == "translate":
        assert not unknown_args
        return run_translate(args, post_dash_dash)
    else:
        executable = "revng-" + command
        if not which(executable, path=search_path):
            log_error("Can't find the following command: {}".format(executable))
            return 1

        in_path = get_command(executable)
        if in_path:
            arguments = [os.path.abspath(in_path)] + all_args
            run(arguments)
        else:
            log_error("\"revng-{}\" is not a valid command".format(command))
            return 1


if __name__ == "__main__":
    sys.exit(main())
