#!/usr/bin/python -O

import re

from os import stat
from os.path import dirname, join
from subprocess import call, PIPE, Popen
from sys import argv, exit, path, stderr


########################################################################


ARGV = argv[1:]

VERBOSE = False
PREPROCESSED = False
PREPROCESSING = False

INPUTS = []
INPUT = None
INSTRUMENTOR_INPUT = None
INSTRUMENTOR_OUTPUT = None

GCC_CC1 = Popen(['/usr/bin/gcc', '-print-prog-name=cc1'], stdout=PIPE).communicate()[0].rstrip()

CPP0 = [GCC_CC1, '-E', '-quiet']
CPP0_OUTPUT_SLOT = None

CC1 = [GCC_CC1, '-fpreprocessed']
CC1_INPUT_SLOT = None

INST = [join(dirname(__file__), 'main')]


def slurp_0(*args):
    global CPP0, CC1
    CPP0 += args
    CC1 += args


def slurp_1(arg):
    global ARGV
    slurp_0(arg, ARGV.pop(0))


def verbose(arg):
    slurp_0(arg)
    global INST, VERBOSE
    INST.append('--show-phase-times')
    VERBOSE = True


def dumpbase(arg):
    global ARGV, CC1, INST
    value = ARGV.pop(0)
    CC1 += [arg, value]
    INST += ['--dumpbase', value]


def inst_toggle(arg):
    if not arg.startswith('-f'):
        print >>stderr, 'malformed instrumentor flag', arg
        exit(1)

    global INST
    INST.append('--' + arg[2:])


def inst_embed(arg):
    global ARGV, INST
    inst_toggle(arg)
    INST.append(ARGV.pop(0))


def ignore(arg):
    pass


def instrumentor_input(arg):
    __pychecker__ = 'no-argsused'
    global ARGV, INSTRUMENTOR_INPUT
    INSTRUMENTOR_INPUT = ARGV.pop(0)


def instrumentor_output(arg):
    __pychecker__ = 'no-argsused'
    global ARGV, INSTRUMENTOR_OUTPUT
    INSTRUMENTOR_OUTPUT = ARGV.pop(0)


def preprocessing(arg):
    __pychecker__ = 'no-argsused'
    global PREPROCESSING
    PREPROCESSING = True


def preprocessed(arg):
    __pychecker__ = 'no-argsused'
    global PREPROCESSED
    PREPROCESSED = True


def cpp0_output_slot(arg):
    global CPP0, CPP0_OUTPUT_SLOT
    slurp_1(arg)
    CPP0_OUTPUT_SLOT = len(CPP0) - 1


def slurp_0_cc1(arg):
    global CC1
    CC1.append(arg)


def slurp_1_cc1(arg):
    global ARGV, CC1
    CC1 += [arg, ARGV.pop(0)]


def slurp_0_cpp0(arg):
    global CPP0
    CPP0.append(arg)


def slurp_1_cpp0(arg):
    global ARGV, CPP0
    CPP0 += [arg, ARGV.pop(0)]


def scheme(name):
    global INST
    INST.append('--scheme-' + name)


def scales(filename):
    global INST
    INST += ['--load-scales', filename]


def clude(flag, arg):
    global INST
    INST += ['--' + flag, arg]


def resetPoints(filename):
    global INST
    INST += ['--reset-at-points', filename]


def skip():
    pass


def cc1_input_slot(arg):
    global CC1, CC1_INPUT_SLOT, INPUTS
    INPUTS.append(arg)
    slurp_0(arg)
    CC1_INPUT_SLOT = len(CC1) - 1


def unhandled(arg):
    print >>stderr, 'warning: unhandled option:', arg
    slurp_0(arg)


FLAG_SPECS_EXACT = {
    # Sampler Options
    '-fsampler-scheme=nothing': ignore,
    '-fthreads': inst_toggle,
    '-fno-threads': inst_toggle,
    '-fsample': inst_toggle,
    '-fno-sample': inst_toggle,
    '-frename-locals': inst_toggle,
    '-fno-rename-locals': inst_toggle,
    '-fassign-across-pointer': inst_toggle,
    '-fno-assign-across-pointer': inst_toggle,
    '-fassign-into-field': inst_toggle,
    '-fno-assign-into-field': inst_toggle,
    '-fassign-into-index': inst_toggle,
    '-fno-assign-into-index': inst_toggle,
    '-fcompare-constants': inst_toggle,
    '-fno-compare-constants': inst_toggle,
    '-ftimestamp-first': inst_toggle,
    '-fno-timestamp-first': inst_toggle,
    '-ftimestamp-last': inst_toggle,
    '-fno-timestamp-last': inst_toggle,
    '-frelative-paths': inst_toggle,
    '-fno-relative-paths': inst_toggle,
    '-fbalance-paths': inst_toggle,
    '-fno-balance-paths': inst_toggle,
    '-fcache-countdown': inst_toggle,
    '-fno-cache-countdown': inst_toggle,
    '-fpredict-checks': inst_toggle,
    '-fno-predict-checks': inst_toggle,
    '-fspecialize-empty-regions': inst_toggle,
    '-fno-specialize-empty-regions': inst_toggle,
    '-fspecialize-singleton-regions': inst_toggle,
    '-fno-specialize-singleton-regions': inst_toggle,
    '-fassume-weighty-externs': inst_toggle,
    '-fno-assume-weighty-externs': inst_toggle,
    '-fassume-weighty-interns': inst_toggle,
    '-fno-assume-weighty-interns': inst_toggle,
    '-fassume-weighty-libraries': inst_toggle,
    '-fno-assume-weighty-libraries': inst_toggle,
    '-fuse-points-to': inst_toggle,
    '-fno-use-points-to': inst_toggle,
    '-fcompare-uninitialized': inst_toggle,
    '-fno-compare-uninitialized': inst_toggle,
    '-fadd-blast-markers': inst_toggle,
    '-fno-add-blast-markers': inst_toggle,
    '-fsave-blast-spec': inst_embed,
    '-fshow-stats': inst_toggle,
    '-fno-show-stats': inst_toggle,
    '-fsave-cfg': inst_embed,
    '-fsave-dataflow': inst_embed,
    '-fsave-dataflow-fields': inst_toggle,
    '-fno-save-dataflow-fields': inst_toggle,
    '-fisolate-shared-accesses': inst_toggle,
    '-fno-isolate-shared-accesses': inst_toggle,
    '-fsave-site-info': inst_embed,
    '-fsave-implications': inst_embed,
    '-fsampler-dataflow': ignore,
    '-fsampler-implications': ignore,
    '-finstrumentor-input': instrumentor_input,
    '-finstrumentor-output': instrumentor_output,
    '-fannotate-sites': inst_toggle,
    '-fextra-edge-profiles': inst_toggle,

     # Overall Options
    '-E': preprocessing,
    '-o': cpp0_output_slot,
    '-v': verbose,
    
    # C Language Options
    '-aux-info': slurp_0_cc1,
    
    # Optimization Options
    '--param': slurp_1_cc1,
    
    # Preprocessor Options
    '-A': slurp_1,
    '-D': slurp_1,

    '-MT': slurp_1_cpp0,
    '-MQ': slurp_1_cpp0,
    '-MP': slurp_0_cpp0,

    '-MF': slurp_1_cpp0,
    '-MD': slurp_1_cpp0,
    '-MMD': slurp_1_cpp0,
    '-I': slurp_1,
    '-idirafter': slurp_1,
    '-include': slurp_1,
    '-imacros': slurp_1,
    '-iprefix': slurp_1,
    '-iwithprefix': slurp_1,
    '-iwithprefixbefore': slurp_1,
    '-isystem': slurp_1,
    '-fpreprocessed': preprocessed,
    '-U': slurp_1,

    # Undocumented Options
    '-auxbase': slurp_1,
    '-auxbase-strip': slurp_1,
    '-dumpbase': dumpbase,
    '-quiet': slurp_0_cc1,
    }


FLAG_SPECS_REGEXP = [
    (re.compile(pattern), handler)
    for pattern, handler in [
        (r'^(-O\d*)$', slurp_0_cc1),
        (r'^(-g\d*)$', slurp_0_cc1),
        (r'^(-m(?:arch|cpu)=.*)$', slurp_0_cc1),
        (r'^-fsampler-scheme=(.*)$', scheme),
        (r'^-fsampler-scales=(.*)$', scales),
        (r'^-f((?:in|ex)clude-(?:file|function|location))=(.*)$', clude),
        (r'^-fsampler-random=', skip),
        (r'^-fsampler-lib-dir=', skip),
        (r'^-freset-at-points=(.*)', resetPoints),
        (r'^(-.+)$', slurp_0),
        (r'^(.*\.[ci]|-)$', cc1_input_slot),
        (r'^(.*)$', unhandled),
        ]]


########################################################################


def run_stage(args, **kwargs):
    global VERBOSE
    if VERBOSE:
        print >>stderr, '  run_stage: %s\n' % ' '.join(args)

    status = call(args, **kwargs)
    if status:
        exit(status)


def main():
    global ARGV
    while ARGV:
        arg = ARGV.pop(0)
        global FLAG_SPECS_EXACT
        handler = FLAG_SPECS_EXACT.get(arg)
        if handler:
            handler(arg)
        else:
            global FLAG_SPECS_REGEXP
            for pattern, handler in FLAG_SPECS_REGEXP:
                match = pattern.match(arg)
                if match:
                    handler(*match.groups())
                    break

    global CPP0, INPUTS, PREPROCESSING
    if PREPROCESSING:
        run_stage(CPP0)
        exit(0)
    elif len(INPUTS) == 1:
        [inarg] = INPUTS
    else:
        print >>stderr, 'expected one input file, but got:', INPUTS
        print >>stderr, '  cc1 command line was probably parsed incorrectly'
        exit(1)

    global INSTRUMENTOR_OUTPUT
    if INSTRUMENTOR_OUTPUT is None:
        print >>stderr, 'internal error: instrumentor output file was not set'
        exit(1)

    global INSTRUMENTOR_INPUT, PREPROCESSED
    if PREPROCESSED:
        INSTRUMENTOR_INPUT = inarg
    else:
        if INSTRUMENTOR_INPUT is None:
            print >>stderr, 'internal error: instrumentor input file was not set'
            exit(1)
        global CPP0_OUTPUT_SLOT
        if CPP0_OUTPUT_SLOT is not None:
            CPP0[CPP0_OUTPUT_SLOT] = INSTRUMENTOR_INPUT
        else:
            CPP0 += ['-o', INSTRUMENTOR_INPUT]
        run_stage(CPP0)

    global INST
    INST.append(INSTRUMENTOR_INPUT)
    run_stage(INST, stdout=open(INSTRUMENTOR_OUTPUT, 'w'))

    if stat(INSTRUMENTOR_OUTPUT).st_size == 0:
        print >>stderr, 'instrumentor produced no output'
        exit(1)

    global CC1, CC1_INPUT_SLOT
    CC1[CC1_INPUT_SLOT] = INSTRUMENTOR_OUTPUT
    run_stage(CC1)


if __name__ == '__main__':
    main()
