#!/usr/bin/python -O

import cPickle as pickle
import os
import re

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


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


ARGV = argv[1:]

CONFIG = pickle.loads(os.environ['SAMPLER_CC_PARSED_CONFIG'])

VERBOSE = CONFIG.verbose
PREPROCESSED = False
PREPROCESSING = False

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

GCC = CONFIG.gcc
GCC_CC1 = Popen([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')]

TOGGLE_PREFIX = 'SAMPLER_CC_TOGGLE_'


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


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


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


def inst_save_blast_spec(arg):
    global ARGV
    filename = ARGV.pop(0)

    add = CONFIG.toggles.get('add-blast-markers')
    if add and int(add):
        global INST
        INST += ['--add-blast-markers', '--save-blast-spec', filename]


def inst_save_cfg(arg):
    global ARGV, INST
    filename = ARGV.pop(0)
    INST += ['--save-cfg', filename]


def inst_save_dataflow(arg):
    global ARGV
    filename = ARGV.pop(0)

    if CONFIG.dataflow:
        global INST
        INST += ['--save-dataflow', filename]


def inst_save_implications(arg):
    global ARGV
    filename = ARGV.pop(0)

    if CONFIG.implications:
        global INST
        INST += ['--save-implications', filename]


def inst_save_site_info(arg):
    global ARGV, INST
    filename = ARGV.pop(0)
    INST += ['--save-site-info', filename]


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 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
    '-finstrumentor-input': instrumentor_input,
    '-finstrumentor-output': instrumentor_output,
    '-fsave-blast-spec': inst_save_blast_spec,
    '-fsave-cfg': inst_save_cfg,
    '-fsave-dataflow': inst_save_dataflow,
    '-fsave-implications': inst_save_implications,
    '-fsave-site-info': inst_save_site_info,

     # Overall Options
    '-E': preprocessing,
    '-o': cpp0_output_slot,
    
    # 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,
    '-imultiarch': slurp_1,
    '-quiet': slurp_0_cc1,
    }


FLAG_SPECS_REGEXP = [
    (re.compile(pattern), handler)
    for pattern, handler in [
        (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 INST
    for aspect, argument in CONFIG.cludes:
        INST += ['--' + aspect, argument]

    if CONFIG.resetAtPoints:
        INST += ['--reset-at-points', CONFIG.resetAtPoints]

    if CONFIG.scales:
        INST += ['--load-scales', CONFIG.scales]

    for scheme in CONFIG.schemes:
        INST.append('--scheme-' + scheme)

    for aspect, enabled in CONFIG.toggles.iteritems():
        negate = ('no-', '')[enabled]
        toggle = '--' + negate + aspect
        INST.append(toggle)

    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)

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

    if os.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()
