#! /usr/bin/env python
#
# Command-line compiler driver for HILTI programs.

import site
import sys
import os.path
import optparse
import subprocess
import time

Version = 0.2

Target = "<unknown>"

Cc = "clang -gfull"               # Note, we really need clang.
CcOptSpec = "-O2"                 # With clang, -O4 is LTO, which we can't do wo/ support by the system linker.
OptOptSpec = "-O2 -std-link-opts"

# Search file in shell PATH. Returns None if not found.
def findInPATH(file):
    hc = os.path.join(os.path.dirname(sys.argv[0]), "../build/tools")
    hc = os.path.abspath(hc)

    for dir in os.environ["PATH"].split(":") + [hc]:
        p = os.path.join(dir, file)
        if os.path.exists(p):
            return p
    else:
        return None

# Find a *-configure tool. Aborts if not found.
def findConfigTool(tool):
    # Look for an environment variable.
    env = tool.upper().replace("-", "_")
    if env in os.environ:
        path = os.environ[env]
        if not os.path.exists(path):
            print >>sys.stderr, "%s set to %s but path does not exist" % (env, path)
            sys.exit(1)

        return os.environ[env]

    # Search in the same directory as we are in.
    dir = os.path.dirname(sys.argv[0])
    path = os.path.join(dir, tool)

    if os.path.exists(path):
        return path

    # Search PATH.
    path = findInPATH(tool)

    if not path:
        print >>sys.stderr, "cannot find %s, aborting." % tool
        sys.exit(1)

    return path

# Run a *-config tool to get an option.
def runConfig(path, option):
    try:
        options = option.split()
        if Options and Options.debug > 0:
            options = ["--debug"] + options

        return subprocess.Popen([path] + options, stdout=subprocess.PIPE).communicate()[0].strip()
    except:
        return ""

# Report a message if verbose output is requested.
def verbose(msg):
    if Options.verbose:
        print >>sys.stderr, msg

def error(msg):
    print >>sys.stderr, msg
    removeTmps()
    sys.exit(1)

# Start timer to measure duration of a tool execution. In verbose mode, it also
# prints the "phase" string given. If time measurement are disabled, it does
# nothing else.
def timeStart(phase):
    if Options.verbose:
        print >>sys.stderr, "  >", phase,

        if Options.time:
            global TimerStart
            TimerStart = time.time()

        else:
            print >>sys.stderr

# Stop timer to measure duration of a tool execution. Does nothing if time
# measurement disabled.
def timeStop():
    if Options.time:
        print >>sys.stderr, "(%.2fs)" % (time.time() - TimerStart)

Tmps = []

# Record a file as a tmp that will be deleted automatically at termination.
def registerTmp(tmp):
    global Tmps
    Tmps += [tmp]

# Deletes all files recorded as being tmp files.
def removeTmps():
    if Options.tmps:
        return

    for output in Tmps:
        try:
            os.unlink(output)
        except OSError:
            pass

# Takes a list of argumens and turns them into space-separated string, but
# make sure to include each only once.
def uniqueArgs(args):
    t = []
    [t.append(i) for i in args if not t.count(i)]
    return " ".join(t)

Used = set()

def makeOutput(input, dstext):
    global used

    cnt = 0
    while True:
        cnt += 1
        (root, ext) = os.path.splitext(input)
        cntext = "-%d" % cnt if cnt > 1 else ""
        output = "%s.hb%s%s.tmp.%s" % (os.path.basename(root), os.getpid(), cntext, dstext)

        if output not in Used:
            break

    registerTmp(output)
    Used.add(output)
    return output

def execute(cmd):
    timeStart(cmd)
    result = os.system(cmd)
    timeStop()
    return result == 0

### Compile and lin function.

# Compile Spicy into HILTI code.
def compilePac(input, output):
    try:
        flags = os.environ.get("SPICYFLAGS" % tag.upper(), "").split()
    except:
        flags = []

    for i in range(Options.debug):
        flags += ["-d"]

    for i in ImportPaths:
        flags += ["-I %s" % i]

    flags = uniqueArgs(flags)

    path = runConfig(SpicyConfig, "--spicyc-binary")

    spicyc = "%s %s %s -o %s" % (path, flags, input, output)

    if not execute(spicyc):
        error("error compiling %s, aborting" % input)

# Compile C and C++ files into LLVM bitcode.
def compileC(input, output, cxx):

    tag = "cxx" if cxx else "c"

    try:
        flags = os.environ.get("%sFLAGS" % tag.upper(), "").split()
    except:
        flags = []

    flags += ["-c", "-emit-llvm", "-g", "-fno-color-diagnostics", "-O0", "-I ."]
    flags += runConfig(HiltiConfig, "--compiler --%sflags" % tag).split()
    flags += runConfig(HiltiConfig, "--runtime --%sflags" % tag).split()

    if Options.spicy:
        flags += runConfig(SpicyConfig, "--compiler --%sflags" % tag).split()
        flags += runConfig(SpicyConfig, "--runtime --%sflags" % tag).split()

    flags = uniqueArgs(flags)

    cc = "%s %s -o %s %s" % (Cc, flags, output, input)

    if not execute(cc):
        error("error compiling %s, aborting" % input)

# Compile HILTI and LLVM files.
def runHiltic(inputs, output, proto=False):
    try:
        flags = os.environ.get("HILTICFLAGS", "").split()
    except:
        flags = []

    if not proto:
        flags += ["-b"]
        # flags += ["-s"]
    else:
        flags += ["-P"]

    if Options.debug:
        # Will add below.
        flags = [f for f in flags if f != "-d"]

    if Options.profile:
        # Will add below.
        flags = [f for f in flags if f != "-F"]

    includes = runConfig(HiltiConfig, "--hilti-libdirs")

    if Options.spicy:
        includes = runConfig(SpicyConfig, "--spicyc-libdirs")

    flags += ["-I%s" % i for i in includes.split()]

    flags += ["-I%s" % i for i in ImportPaths]

    flags = uniqueArgs(flags)

    for i in range(Options.debug):
        flags += " -d"

    for i in range(Options.profile):
        flags += " -F"

    inputs = " ".join(inputs)
    path = runConfig(HiltiConfig, "--hiltic-binary")

    if proto:
        path = os.path.join(os.path.dirname(path), "hiltip")

    hiltic = "%s %s %s -o %s" % (path, inputs, flags, output)

    if not execute(hiltic):
        error("error running hiltic, aborting")

def optimize(input, output, strip_debug=False):
    opt = "opt -f %s -o=%s %s" % (OptOptSpec, output, input)
    return execute(opt)

# Not used currently.
# def stripDebug(input, output):
#     opt = "opt -strip-debug -o=%s %s" % (output, input)
#     return execute(opt)

def nativeLink(input, output):
    try:
        flags = os.environ.get("LDFLAGS", "").split()
    except:
        flags = []

    flags += ["-g"]
    flags += runConfig(HiltiConfig, "--runtime --ldflags").split()
    flags += runConfig(HiltiConfig, "--runtime --libs").split()

    if Options.spicy:
        flags += runConfig(SpicyConfig, "--runtime --ldflags").split()
        flags += runConfig(SpicyConfig, "--runtime --libs").split()

    if Options.optimize:
        flags += [CcOptSpec]

    flags = uniqueArgs(flags)

    cc = "%s %s -o %s %s" % (Cc, flags, output, input)

    if not execute(cc):
        error("error linking %s to native executable, aborting" % input)

### Command line options.

def parseOptions():
    # Collect additional -I arguments.
    def import_path_callback(option, opt, value, parser):
        global ImportPaths
        ImportPaths += [value]

    version = "hilti-build %s\ntarget: %s" % (Version, Target)
    usage = """%prog [options] -o <output> <inputs>

Input files may be end in one of the following extensions:

    *.hlt     - HILTI files
    *.ll      - LLVM assember files
    *.bc      - LLVM bitcode files
    *.c       - C files
    *.cc      - C++ files
    *.spicy    - Spicy files
"""

    optparser = optparse.OptionParser(usage="%prog [options] <input-file>", version=version)

    optparser.add_option("-o", "--output", action="store", type="string", dest="output", default=None,
                         help="Name of destination (mandatory).", metavar="FILE")
    optparser.add_option("-T", "--time", action="store_true", dest="time", default=False,
                         help="Time individual steps (implies --verbose)")
    optparser.add_option("-t", "--tmps", action="store_true", dest="tmps", default=False,
                         help="Create and leave files with intermediary outputs.")
    optparser.add_option("-l", "--llvm", action="store_true", dest="print_llvm", default=False,
                         help="Write out the final LLVM assembly code.")
    optparser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False,
                         help="Output information about intermediary steps.")
    optparser.add_option("-O", "--optimize", action="store_true", dest="optimize", default=False,
                         help="Apply optimizations.")
    optparser.add_option("-P", "--prototypes", action="store_true", dest="proto", default=False,
                         help="Generate just prototypes for all *.hlt files (including generated ones).")
    optparser.add_option("-b", "--bc", action="store_true", dest="bitcode", default=False,
                         help="Produce a bitcode file for JITing rather than an executable.")
    optparser.add_option("-d", "--debug", action="count", dest="debug", default=0,
                         help="Compile HILTI code with debugging support; multiple times increases level.")
    optparser.add_option("-B", "--spicy", action="store_true", dest="spicy", default=False,
                         help="Activate Spicy support even if not *.spicy files are given.")
#    optparser.add_option("-S", "--stack-size", action="store", type="int", dest="stack", default=0,
#                         help="Default HILTI stack size. Default is allocate each frame independently.")
    optparser.add_option("-I", "--import-path", action="callback", callback=import_path_callback, type="string",
                         help="Add DIR to search path for imports", metavar="DIR")
    optparser.add_option("-F", "--profile", action="count", dest="profile", default=0,
                         help="Enable profiling support; each time this option is given, the profiling level is increased by one")

    addl = os.environ.get("HILTI_BUILD_FLAGS", "").split()

    (options, args) = optparser.parse_args(sys.argv[1:] + addl)

    if len(args) < 1:
        optparser.error("no input file specified")

    if not options.output and not options.print_llvm and not options.proto:
        optparser.error("no output file specified")

    if options.time:
        options.verbose = True

    if options.optimize:
        options.debug = 0

    global Options
    Options = options

    return args

### Main.

HiltiConfig = findConfigTool("hilti-config")
SpicyConfig = findConfigTool("spicy-config")
ImportPaths = []

inputs = parseOptions()

files = {}

for ext in [".hlt", ".c", ".cc", ".cxx", ".cpp", ".ll", ".bc", ".spicy"]:
    files[ext] = []

for input in inputs:

    (root, ext) = os.path.splitext(input)

    if ext in files:
        files[ext] += [input]

    else:
        error("unsupported input file %s" % input)

if not Options.spicy:
    Options.spicy = (len(files[".spicy"]) != 0)

if Options.proto:
    for input in files[".hlt"]:
        output = Options.output if Options.output else os.path.basename(input) + ".h"
        runHiltic([input], output, proto=True)

    sys.exit(0)

for input in files[".spicy"]:
    verbose("compiling %s ..." % input)
    output = makeOutput(input, "hlt")
    compilePac(input, output)
    files[".hlt"] += [output]

for input in files[".c"]:
    verbose("compiling %s ..." % input)
    output = makeOutput(input, "bc")
    compileC(input, output, False)
    files[".bc"] += [output]

for input in files[".cc"] + files[".cxx"] + files[".cpp"]:
    verbose("compiling %s ..." % input)
    output = makeOutput(input, "bc")
    compileC(input, output, True)
    files[".bc"] += [output]

output = Options.output if Options.output else "no-output-name"
output = makeOutput(output, "bc")

runHiltic(files[".hlt"] + files[".ll"] + files[".bc"], output)

if Options.optimize:
    input = output
    output = makeOutput(output, "opt.bc")
    optimize(input, output)

input = output

if Options.print_llvm:
    output = ">%s" % Options.output if Options.output else ""
    execute("llvm-dis <%s %s" % (input, output))

elif Options.bitcode:
    # We're done.
    os.rename(input, Options.output)

else:
    # Produce the executable.
    nativeLink(input, Options.output)

removeTmps()
sys.exit(0)
