import lit.formats
import lit.util
import lit
import os, glob, re
import shlex
import pipes
from lit.formats import FileBasedTest
from lit.TestRunner import executeScript, executeScriptInternal, parseIntegratedTestScriptCommands, getDefaultSubstitutions, applySubstitutions, getTempPaths
from lit import Test
from lit.util import to_bytes, to_string
try:
    from shlex import quote # python 3.3 and above
except:
    from pipes import quote # python 3.2 and earlier

def parseBenchmarkScript(test):
    """Scan a llvm-testsuite like benchmark .test script."""
    def parseShellCommand(script, ln):
        # Trim trailing whitespace.
        ln = ln.rstrip()

        # Collapse lines with trailing '\\'.
        if script and script[-1][-1] == '\\':
            script[-1] = script[-1][:-1] + ln
        else:
            script.append(ln)

    # Collect the test lines from the script.
    sourcepath = test.getSourcePath()
    runscript = []
    verifyscript = []
    keywords = ['RUN:', 'VERIFY:']
    for line_number, command_type, ln in \
            parseIntegratedTestScriptCommands(sourcepath, keywords):
        if command_type == 'RUN':
            parseShellCommand(runscript, ln)
        elif command_type == 'VERIFY':
            parseShellCommand(verifyscript, ln)
        else:
            raise ValueError("unknown script command type: %r" % (
                    command_type,))

    # Verify the script contains a run line.
    if runscript == []:
        return lit.Test.Result(Test.UNRESOLVED, "Test has no RUN: line!")

    # Check for unterminated run lines.
    for script in runscript, verifyscript:
        if script and script[-1][-1] == '\\':
            return lit.Test.Result(Test.UNRESOLVED,
                                   "Test has unterminated RUN/VERIFY lines (with '\\')")

    return runscript,verifyscript

def getUserTimeFromTimeOutput(f):
    with open(f) as fd:
        l = [l for l in fd.readlines()
             if l.startswith('user')]
    assert len(l) == 1

    m = re.match(r'user\s+([0-9.]+)', l[0])
    return float(m.group(1))

def collectCompileTime(test):
    # TODO: This is not correct yet as the directory may contain .o.time files
    # of multiple benchmarks in the case of SingleSource tests.
    compile_time = 0.0
    basepath = os.path.dirname(test.getFilePath())
    for path, subdirs, files in os.walk(basepath):
        for file in files:
            if file.endswith('.o.time'):
                compile_time += getUserTimeFromTimeOutput(os.path.join(path, file))
    return compile_time

def runScript(test, litConfig, script, tmpBase, useExternalSh = True):
    execdir = os.path.dirname(test.getExecPath())
    if useExternalSh:
        res = executeScript(test, litConfig, tmpBase, script, execdir)
    else:
        res = executeScriptInternal(test, litConfig, tmpBase, script, execdir)
    return res

def prepareRunSafely(config, commandline, outfile):
    stdin = None
    stdout = None
    stderr = None
    workdir = None
    tokens = shlex.split(commandline)
    # Parse "< INPUTFILE", "> OUTFILE", "2> OUTFILE" patterns
    i = 0
    while i < len(tokens):
        if tokens[i] == "<" and i+1 < len(tokens):
            stdin = tokens[i+1]
            del tokens[i+1]
            del tokens[i]
            continue
        elif tokens[i] == ">" and i+1 < len(tokens):
            stdout = tokens[i+1]
            del tokens[i+1]
            del tokens[i]
            continue
        elif tokens[i] == "2>" and i+1 < len(tokens):
            stderr = tokens[i+1]
            del tokens[i+1]
            del tokens[i]
            continue
        if i+2 < len(tokens) and tokens[i] == "cd" and tokens[i+2] == ";":
            workdir = tokens[i+1]
            del tokens[i+2]
            del tokens[i+1]
            del tokens[i]
            continue
        i += 1

    runsafely = "%s/RunSafely.sh" % config.test_suite_root
    runsafely_prefix = [ runsafely ]
    if workdir is not None:
        runsafely_prefix += [ "-d", workdir ]
    timeit = "%s/tools/timeit" % config.test_source_root
    if config.remote_host:
        timeit = "%s/tools/timeit-target" % config.test_source_root
        runsafely_prefix += [ "-r", config.remote_host ]
        if config.remote_user:
            runsafely_prefix += [ "-l", config.remote_user ]
        if config.remote_client:
            runsafely_prefix += [ "-rc", config.remote_client ]
        if config.remote_port:
            runsafely_prefix += [ "-rp", config.remote_port ]
    if config.run_under:
        runsafely_prefix += [ "-u", config.run_under ]
    if not config.traditional_output:
        runsafely_prefix += [ "-n" ]
        if stdout is not None:
            runsafely_prefix += [ "-o", stdout ]
        if stderr is not None:
            runsafely_prefix += [ "-e", stderr ]
    else:
        if stdout is not None or stderr is not None:
            raise Exception("separate stdout/stderr redirection not possible with traditional output")
    timeout = "7200"
    if stdin is None:
        stdin = "/dev/null"
    runsafely_prefix += [ "-t", timeit, timeout, stdin, outfile ]

    new_commandline = " ".join(map(quote, runsafely_prefix + tokens))
    return new_commandline

class TestSuiteTest(FileBasedTest):
    def __init__(self):
        super(TestSuiteTest, self).__init__()
        
    def execute(self, test, litConfig):
        config = test.config
        if config.unsupported:
            return (Test.UNSUPPORTED, 'Test is unsupported')

        # Parse benchmark script
        res = parseBenchmarkScript(test)
        if isinstance(res, lit.Test.Result):
            return res
        if litConfig.noExecute:
            return lit.Test.Result(Test.PASS)
        runscript, verifyscript = res

        tmpDir, tmpBase = getTempPaths(test)
        outfile = tmpBase + ".out"
        substitutions = getDefaultSubstitutions(test, tmpDir, tmpBase)
        substitutions += [('%o', outfile)]
        runscript = applySubstitutions(runscript, substitutions)
        verifyscript = applySubstitutions(verifyscript, substitutions)

        adjusted_runscript = []
        for line in runscript:
            line = prepareRunSafely(config, line, outfile)
            adjusted_runscript.append(line)
        runscript = adjusted_runscript

        # Create the output directory if it does not already exist.
        lit.util.mkdir_p(os.path.dirname(tmpBase))

        # Execute runscript (the "RUN:" part)
        output = ""
        n_runs = 1
        runtimes = []
        for n in range(n_runs):
            res = runScript(test, litConfig, runscript, tmpBase)
            if isinstance(res, lit.Test.Result):
                return res

            output += "\n" + "\n".join(runscript)

            out, err, exitCode, timeoutInfo = res
            if exitCode == Test.FAIL:
                # Only show command output in case of errors
                output += "\n" + out
                output += "\n" + err
                return lit.Test.Result(Test.FAIL, output)

            timefile = "%s.time" % (outfile,)
            try:
                runtime = getUserTimeFromTimeOutput(timefile)
                runtimes.append(runtime)
            except IOError:
                pass

        # Run verification script (the "VERIFY:" part)
        if len(verifyscript) > 0:
            res = runScript(test, litConfig, verifyscript, tmpBase)
            if isinstance(res, lit.Test.Result):
                return res
            out, err, exitCode, timeoutInfo = res

            output += "\n" + "\n".join(verifyscript)
            if exitCode != 0:
                output += "\n" + out
                output += "\n" + err
                return lit.Test.Result(Test.FAIL, output)

        # Put metrics into the test result.
        result = lit.Test.Result(Test.PASS, output)
        if len(runtimes) > 0:
            result.addMetric('exec_time', lit.Test.toMetricValue(runtimes[0]))
        try:
            compile_time = collectCompileTime(test)
            result.addMetric('compile_time',
                             lit.Test.toMetricValue(compile_time))
        except IOError:
            pass

        return result

config.name = 'test-suite'
config.test_format = TestSuiteTest()
config.suffixes = ['.test']
config.excludes = ['ABI-Testsuite']
config.traditional_output = True
if 'SSH_AUTH_SOCK' in os.environ:
    config.environment['SSH_AUTH_SOCK'] = os.environ['SSH_AUTH_SOCK']