#!/usr/bin/env python3

'''
The fuzzer tool that transforms existing problems.
'''

import sys
import argparse
import random

from stringfuzz.constants import LANGUAGES, SMT_20_STRING, SMT_25_STRING
from stringfuzz.transformers import unprintable, nop, rotate, fuzz, graft, translate, reverse, multiply
from stringfuzz.generator import generate
from stringfuzz.parser import parse, ParsingError
from stringfuzz.ast import SettingNode, ExpressionNode, MetaCommandNode

# constants
UNPRINTABLE = 'unprintable'
NOP         = 'nop'
ROTATE      = 'rotate'
FUZZ        = 'fuzz'
GRAFT       = 'graft'
TRANSLATE   = 'translate'
REVERSE     = 'reverse'
MULTIPLY    = 'multiply'

TRANSFORMERS = {
    UNPRINTABLE: unprintable,
    NOP:         nop,
    ROTATE:      rotate,
    FUZZ:        fuzz,
    GRAFT:       graft,
    TRANSLATE:   translate,
    REVERSE:     reverse,
    MULTIPLY:    multiply
}

# defaults
DEFAULT_SEED           = 0
DEFAULT_RANDOM         = False
DEFAULT_FACTOR         = 2
DEFAULT_INTEGER_FLAG   = False
DEFAULT_SKIP_RE_RANGE  = True
DEFAULT_SKIP_STR_TO_RE = True

GET_MODEL = "get-model"
GET_INFO  = "get-info"
TO_STRIP  = [GET_MODEL, GET_INFO]

def should_keep(expr):
    if isinstance(expr, SettingNode):
        return False
    if isinstance(expr, MetaCommandNode):
        return False
    if isinstance(expr, ExpressionNode):
        if expr.symbol in TO_STRIP:
            return False
    return True

def main():

    # create arg parser
    global_parser = argparse.ArgumentParser(description='SMTLIB 2.* problem transformer.')
    global_parser.add_argument(
        '--file',
        '-f',
        dest    = 'input_file',
        metavar = 'F',
        default = sys.stdin,
        type    = argparse.FileType('r'),
        help    = 'input file (default: stdin)'
    )
    global_parser.add_argument(
        '--in-lang',
        '-i',
        dest    = 'input_language',
        type    = str,
        choices = LANGUAGES,
        default = SMT_25_STRING,
        help    = 'input language (default: {})'.format(SMT_25_STRING)
    )
    global_parser.add_argument(
        '--out-lang',
        '-o',
        dest    = 'output_language',
        type    = str,
        choices = LANGUAGES,
        default = SMT_25_STRING,
        help    = 'output language (default: {})'.format(SMT_25_STRING)
    )

    seed_group = global_parser.add_mutually_exclusive_group()
    seed_group.add_argument(
        '--seed',
        '-s',
        dest    = 'seed',
        metavar = 'S',
        type    = int,
        default = DEFAULT_SEED,
        help    = 'seed for random number generator (default: {})'.format(DEFAULT_SEED)
    )
    seed_group.add_argument(
        '--random',
        '-r',
        dest    = 'random',
        action  = 'store_true',
        default = DEFAULT_RANDOM,
        help    = 'seed the random number generator with the current time (default: {})'.format(DEFAULT_RANDOM)
    )

    # get subparsers
    subparsers = global_parser.add_subparsers(dest='transformer', help='transformer choice')
    subparsers.required = True

    # fuzz transformer
    fuzz_parser = subparsers.add_parser(FUZZ, help='fuzz transformer')
    fuzz_parser.add_argument(
        '--re-range',
        dest    = 'skip_re_range',
        action  = 'store_false',
        default = DEFAULT_SKIP_RE_RANGE,
        help    = 'Include re_range nodes in multiplication (default: {})'.format(DEFAULT_SKIP_RE_RANGE)
    )
    # graft transformer
    graft_parser = subparsers.add_parser(GRAFT, help='graft transformer')
    graft_parser.add_argument(
        '--str-to-re',
        dest    = 'skip_str_to_re',
        action  = 'store_false',
        default = DEFAULT_SKIP_STR_TO_RE,
        help    = 'Include str_to_re nodes in grafting (default: {})'.format(DEFAULT_SKIP_STR_TO_RE)
    )

    # multiply transformer
    multiply_parser = subparsers.add_parser(MULTIPLY, help='multiply transformer')
    multiply_parser.add_argument(
        '--factor',
        '-n',
        dest    = 'factor',
        metavar = 'N',
        type    = int,
        default = DEFAULT_FACTOR,
        help    = 'number to multiply literals by (default: {})'.format(DEFAULT_FACTOR)
    )
    multiply_parser.add_argument(
        '--re-range',
        dest    = 'skip_re_range',
        action  = 'store_false',
        default = DEFAULT_SKIP_RE_RANGE,
        help    = 'Include re_range nodes in multiplication (default: {})'.format(DEFAULT_SKIP_RE_RANGE)
    )

    # nop transformer
    nop_parser = subparsers.add_parser(NOP, help='nop transformer')

    # reverse transformer
    reverse_parser = subparsers.add_parser(REVERSE, help='reverse transformer')

    # rotate transformer
    rotate_parser = subparsers.add_parser(ROTATE, help='rotate transformer')

    # translate transformer
    translate_parser = subparsers.add_parser(TRANSLATE, help='translate transformer')
    translate_parser.add_argument(
        '--integers',
        dest    = 'translate_integers',
        action  = 'store_true',
        default = DEFAULT_INTEGER_FLAG,
        help    = 'Include integers in translation (default: {})'.format(DEFAULT_INTEGER_FLAG)
    )
    translate_parser.add_argument(
        '--re-range',
        dest    = 'skip_re_range',
        action  = 'store_false',
        default = DEFAULT_SKIP_RE_RANGE,
        help    = 'Include re_range nodes in translation (default: {})'.format(DEFAULT_SKIP_RE_RANGE)
    )

    # unprintable transformer
    unprintable_parser = subparsers.add_parser(UNPRINTABLE, help='unprintable transformer')

    # parse args
    args = global_parser.parse_args()

    # get the transformer function based on args
    transformer_name = args.transformer
    transformer      = TRANSFORMERS[transformer_name]

    # get some flags that will get popped from args before they're used
    input_file      = args.input_file
    input_language  = args.input_language
    output_language = args.output_language

    # seed the RNG
    if args.random is True:
        random.seed()
    else:
        random.seed(args.seed)

    # read input
    raw_in = args.input_file.read()

    # parse input
    try:
        ast = parse(raw_in, input_language)
    except ParsingError as e:
        print(e, file=sys.stderr)
        return 1

    # the nop transformer should not modify anything
    if transformer != nop:

        # filter out suppressed expressions
        ast = list(filter(should_keep, ast))

    # get args as a dict
    transformer_args = vars(args)

    # pop arguments that are specific to this script because
    # they shouldn't be passed on to the transformer
    transformer_args.pop('input_file')
    transformer_args.pop('input_language')
    transformer_args.pop('output_language')
    transformer_args.pop('seed')
    transformer_args.pop('random')
    transformer_args.pop('transformer')

    # run the transformer with the args
    transformed = transformer(ast, **transformer_args)

    # transformers produce ASTs
    print(generate(transformed, output_language))

if __name__ == '__main__':
    main()
