#!/usr/bin/env python3

'''
The fuzzer tool that generates new problems.
'''

import sys
import argparse
import random

from stringfuzz.constants import LANGUAGES, SMT_20_STRING, SMT_25_STRING
from stringfuzz.parser import parse
from stringfuzz.generator import generate
from stringfuzz.smt import smt_get_model, smt_string_logic

from stringfuzz.generators import concats, SYNTACTIC_DEPTH, SEMANTIC_DEPTH
from stringfuzz.generators import overlaps
from stringfuzz.generators import lengths
from stringfuzz.generators import random_text
from stringfuzz.generators import random_ast
from stringfuzz.generators import equality
from stringfuzz.generators import regex, INCREASING_LITERALS, RANDOM_LITERALS, MEMBER_IN, MEMBER_NOT_IN, MEMBER_ALTERNATING, MEMBER_RANDOM, OPERATOR_STAR, OPERATOR_PLUS, OPERATOR_UNION, OPERATOR_INTER, OPERATOR_CONCAT, OPERATOR_ALTERNATING, OPERATOR_RANDOM

# constants
CONCATS     = 'concats'
LENGTHS     = 'lengths'
OVERLAPS    = 'overlaps'
RANDOM_TEXT = 'random-text'
RANDOM_AST  = 'random-ast'
REGEX       = 'regex'
EQUALITY    = 'equality'

GENERATORS = {
    CONCATS:     concats,
    LENGTHS:     lengths,
    OVERLAPS:    overlaps,
    RANDOM_TEXT: random_text,
    RANDOM_AST:  random_ast,
    REGEX:       regex,
    EQUALITY:    equality,
}

DEPTH_TYPES      = [SYNTACTIC_DEPTH, SEMANTIC_DEPTH]
LITERAL_TYPES    = [INCREASING_LITERALS, RANDOM_LITERALS]
MEMBERSHIP_TYPES = [MEMBER_IN, MEMBER_NOT_IN, MEMBER_ALTERNATING, MEMBER_RANDOM]
OPERATOR_LIST    = [OPERATOR_STAR, OPERATOR_PLUS, OPERATOR_UNION, OPERATOR_INTER, OPERATOR_CONCAT]
OPERATOR_TYPES   = [OPERATOR_ALTERNATING, OPERATOR_RANDOM]

# defaults
DEFAULT_SEED           = 0
DEFAULT_RANDOM         = False
DEFAULT_PRODUCE_MODELS = False

DEFAULT_LENGTH           = 10
DEFAULT_DEPTH            = 5
DEFAULT_SOLUTION         = None
DEFAULT_BALANCED         = False
DEFAULT_RANDOM_RELATIONS = False
DEFAULT_DEPTH_TYPE       = SYNTACTIC_DEPTH
DEFAULT_EXTRACTED        = 0
DEFAULT_EXTRACT_INDEX    = 10

DEFAULT_NUM_VARS    = 5
DEFAULT_MIN_LEN     = 0
DEFAULT_MAX_LEN     = 10
DEFAULT_NUM_CONCATS = 0

DEFAULT_NUM_OVERLAPPING  = 2
DEFAULT_LENGTH_OF_CONSTS = 5

DEFAULT_NUM_EQUAL_EXPRESSIONS   = 2
DEFAULT_TERMS_PER_EXPRESSION    = 7
DEFAULT_PREFIX_LENGTH           = 5
DEFAULT_SUFFIX_LENGTH           = 0
DEFAULT_INFIX_LENGTH            = 3
DEFAULT_INFIX_PROBABILITY       = 0.2
DEFAULT_ADD_INFIXES             = False
DEFAULT_EQUAL_RANDOMISE_LENGTHS = False

DEFAULT_LITERAL_TYPE    = RANDOM_LITERALS
DEFAULT_MEMBERSHIP_TYPE = MEMBER_IN
DEFAULT_REGEX_COINCIDE  = False
DEFAULT_NUM_REGEX_TERMS = 3
DEFAULT_NUM_REGEXES     = 1
DEFAULT_REGEX_LIT_MIN   = 1
DEFAULT_REGEX_LIT_MAX   = 3
DEFAULT_TERM_DEPTH      = 0
DEFAULT_REGEX_VAR_MIN   = None
DEFAULT_REGEX_VAR_MAX   = None
DEFAULT_OPERATORS       = OPERATOR_STAR + OPERATOR_PLUS + OPERATOR_UNION
DEFAULT_OPERATOR_TYPE   = OPERATOR_RANDOM

DEFAULT_NUM_RANDOM_VARS     = 5
DEFAULT_NUM_RANDOM_ASSERTS  = 10
DEFAULT_RANDOM_DEPTH        = 2
DEFAULT_MAX_RANDOM_TERMS    = 5
DEFAULT_MAX_RANDOM_STRINGS  = 10
DEFAULT_MAX_RANDOM_NUMBERS  = 10
DEFAULT_SEMANTICALLY_VALID  = False
DEFAULT_LITERAL_PROBABILITY = 0.1

def main():

    # create arg parser
    global_parser = argparse.ArgumentParser(description='SMTLIB 2.* problem generator.')

    # global args
    global_parser.add_argument(
        '--language',
        '-l',
        dest    = 'language',
        type    = str,
        choices = LANGUAGES,
        default = SMT_25_STRING,
        help    = 'output language (default: {})'.format(SMT_25_STRING)
    )
    global_parser.add_argument(
        '--models',
        '-m',
        dest    = 'produce_models',
        action  = 'store_true',
        default = DEFAULT_PRODUCE_MODELS,
        help    = 'append the SMT 2.x command to produce a model (default: {})'.format(DEFAULT_PRODUCE_MODELS)
    )
    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='generator', help='generator choice')
    subparsers.required = True

    # concats fuzzer
    concats_parser = subparsers.add_parser(CONCATS, help='instance with deeply nested concats')
    concats_parser.add_argument(
        '--depth',
        '-d',
        dest    = 'depth',
        metavar = 'D',
        type    = int,
        default = DEFAULT_DEPTH,
        help    = 'depth of the concats (default: {})'.format(DEFAULT_DEPTH)
    )
    concats_parser.add_argument(
        '--depth-type',
        '-t',
        dest    = 'depth_type',
        type    = str,
        choices = DEPTH_TYPES,
        default = DEFAULT_DEPTH_TYPE,
        help    = 'type of depth (default: {})'.format(DEFAULT_DEPTH_TYPE)
    )
    concats_parser.add_argument(
        '--solution',
        '-s',
        dest    = 'solution',
        metavar = 'S',
        type    = str,
        default = DEFAULT_SOLUTION,
        help    = 'expected solution (default: {!r})'.format(DEFAULT_SOLUTION)
    )
    concats_parser.add_argument(
        '--extract',
        '-e',
        dest    = 'num_extracts',
        metavar = 'N',
        type    = int,
        default = DEFAULT_EXTRACTED,
        help    = 'number of extracts to add (default: {})'.format(DEFAULT_EXTRACTED)
    )
    concats_parser.add_argument(
        '--extract-max',
        '-m',
        dest    = 'max_extract_index',
        metavar = 'N',
        type    = int,
        default = DEFAULT_EXTRACT_INDEX,
        help    = 'max index from which to extract (default: {})'.format(DEFAULT_EXTRACT_INDEX)
    )
    concats_parser.add_argument(
        '--balanced',
        '-b',
        action  = 'store_true',
        default = DEFAULT_BALANCED,
        help    = 'flag for balanced tree (default: {!r})'.format(DEFAULT_BALANCED)
    )

    # lengths fuzzer
    lengths_parser = subparsers.add_parser(LENGTHS, help='instance with length constraints')
    lengths_parser.add_argument(
        '--num-vars',
        '-v',
        dest    = 'num_vars',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_VARS,
        help    = 'number of variables to create (default: {})'.format(DEFAULT_NUM_VARS)
    )
    lengths_parser.add_argument(
        '--min-length',
        '-n',
        dest    = 'min_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MIN_LEN,
        help    = 'lower bound on length (default: {})'.format(DEFAULT_MIN_LEN)
    )
    lengths_parser.add_argument(
        '--max-length',
        '-x',
        dest    = 'max_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MAX_LEN,
        help    = 'upper bound on length (default: {})'.format(DEFAULT_MAX_LEN)
    )
    lengths_parser.add_argument(
        '--num-concats',
        '-c',
        dest    = 'num_concats',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_CONCATS,
        help    = 'number of (binary) concats to add (default: {})'.format(DEFAULT_NUM_CONCATS)
    )
    lengths_parser.add_argument(
        '--random-relations',
        '-r',
        action  = 'store_true',
        dest    = 'random_relations',
        default = DEFAULT_RANDOM_RELATIONS,
        help    = 'use constraints other than "=" (default: {!r})'.format(DEFAULT_RANDOM_RELATIONS)
    )

    # overlaps fuzzer
    overlaps_parser = subparsers.add_parser(OVERLAPS, help='instance with overlapping variables')
    overlaps_parser.add_argument(
        '--num-vars',
        '-n',
        dest    = 'num_vars',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_OVERLAPPING,
        help    = 'number of overlapping variables to generate (default: {})'.format(DEFAULT_NUM_OVERLAPPING)
    )
    overlaps_parser.add_argument(
        '--length-of-consts',
        '-c',
        dest    = 'length_of_consts',
        metavar = 'N',
        type    = int,
        default = DEFAULT_LENGTH_OF_CONSTS,
        help    = 'the length of the constant terms (default: {})'.format(DEFAULT_LENGTH_OF_CONSTS)
    )

    # equality fuzzer
    equality_parser = subparsers.add_parser(EQUALITY, help='instance with concatenated expressions (of mixed constants and variables) all equal to each other')
    equality_parser.add_argument(
        '--num-exprs',
        '-n',
        dest    = 'num_expressions',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_EQUAL_EXPRESSIONS,
        help    = 'number of equal expressions to generate (default: {})'.format(DEFAULT_NUM_EQUAL_EXPRESSIONS)
    )
    equality_parser.add_argument(
        '--num-terms',
        '-t',
        dest    = 'num_terms',
        metavar = 'N',
        type    = int,
        default = DEFAULT_TERMS_PER_EXPRESSION,
        help    = 'number terms in each expression (default: {})'.format(DEFAULT_TERMS_PER_EXPRESSION)
    )
    equality_parser.add_argument(
        '--prefix-len',
        '-p',
        dest    = 'prefix_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_PREFIX_LENGTH,
        help    = 'length of constant prefix (default: {})'.format(DEFAULT_PREFIX_LENGTH)
    )
    equality_parser.add_argument(
        '--suffix-len',
        '-s',
        dest    = 'suffix_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_SUFFIX_LENGTH,
        help    = 'length of constant suffix (default: {})'.format(DEFAULT_SUFFIX_LENGTH)
    )
    equality_parser.add_argument(
        '--infix-len',
        '-i',
        dest    = 'infix_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_INFIX_LENGTH,
        help    = 'lengths of constant infixes (default: {})'.format(DEFAULT_INFIX_LENGTH)
    )
    equality_parser.add_argument(
        '--infix-chance',
        '-c',
        dest    = 'infix_probability',
        metavar = 'P',
        type    = float,
        default = DEFAULT_INFIX_PROBABILITY,
        help    = 'probability of constant infixes replacing variables (default: {})'.format(DEFAULT_INFIX_PROBABILITY)
    )
    equality_parser.add_argument(
        '--add-infixes',
        '-m',
        dest    = 'add_infixes',
        action  = 'store_true',
        default = DEFAULT_ADD_INFIXES,
        help    = 'add constant infixes to expressions with 50%% chance (default: {})'.format(DEFAULT_ADD_INFIXES)
    )
    equality_parser.add_argument(
        '--random-lengths',
        '-r',
        dest    = 'randomise_lengths',
        action  = 'store_true',
        default = DEFAULT_EQUAL_RANDOMISE_LENGTHS,
        help    = 'treat all length settings as upper bounds only (default: {})'.format(DEFAULT_EQUAL_RANDOMISE_LENGTHS)
    )

    # regex-pair fuzzer
    regex_parser = subparsers.add_parser(REGEX, help='instance testing one variable for regex membership')
    regex_parser.add_argument(
        '--literal-type',
        '-a',
        dest    = 'literal_type',
        type    = str,
        choices = LITERAL_TYPES,
        default = DEFAULT_LITERAL_TYPE,
        help    = 'way to generate regex literals (default: {})'.format(DEFAULT_LITERAL_TYPE)
    )
    regex_parser.add_argument(
        '--coincide',
        '-c',
        dest    = 'reset_alphabet',
        action  = 'store_true',
        default = DEFAULT_REGEX_COINCIDE,
        help    = 'try to make the regexes share alphabets (default: {})'.format(DEFAULT_REGEX_COINCIDE)
    )
    regex_parser.add_argument(
        '--membership-type',
        '-i',
        dest    = 'membership_type',
        type    = str,
        choices = MEMBERSHIP_TYPES,
        default = DEFAULT_MEMBERSHIP_TYPE,
        help    = 'way to test regex membership (default: {})'.format(DEFAULT_MEMBERSHIP_TYPE)
    )
    regex_parser.add_argument(
        '--num-terms',
        '-t',
        dest    = 'num_terms',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_REGEX_TERMS,
        help    = 'number of terms in each regex (default: {})'.format(DEFAULT_NUM_REGEX_TERMS)
    )
    regex_parser.add_argument(
        '--lit-min',
        '-m',
        dest    = 'literal_min',
        metavar = 'N',
        type    = int,
        default = DEFAULT_REGEX_LIT_MIN,
        help    = 'min length of regex terms (default: {})'.format(DEFAULT_REGEX_LIT_MIN)
    )
    regex_parser.add_argument(
        '--lit-max',
        '-x',
        dest    = 'literal_max',
        metavar = 'N',
        type    = int,
        default = DEFAULT_REGEX_LIT_MAX,
        help    = 'max length of regex terms (default: {})'.format(DEFAULT_REGEX_LIT_MAX)
    )
    regex_parser.add_argument(
        '--depth',
        '-d',
        dest    = 'term_depth',
        metavar = 'D',
        type    = int,
        default = DEFAULT_TERM_DEPTH,
        help    = 'depth of terms (default: {})'.format(DEFAULT_TERM_DEPTH)
    )
    regex_parser.add_argument(
        '--num-regexes',
        '-r',
        dest    = 'num_regexes',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_REGEXES,
        help    = 'number of regexes to test (default: {})'.format(DEFAULT_NUM_REGEXES)
    )
    regex_parser.add_argument(
        '--var-min',
        '-M',
        dest    = 'min_var_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_REGEX_VAR_MIN,
        help    = 'min length of the variable (default: {})'.format(DEFAULT_REGEX_VAR_MIN)
    )
    regex_parser.add_argument(
        '--var-max',
        '-X',
        dest    = 'max_var_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_REGEX_VAR_MAX,
        help    = 'max length of the variable (default: {})'.format(DEFAULT_REGEX_VAR_MAX)
    )
    regex_parser.add_argument(
        '--operators',
        '-o',
        dest    = 'operators',
        metavar = 'OPS',
        type    = str,
        default = DEFAULT_OPERATORS,
        help    = 'operators to choose from for deep terms. OPS is of the form [{}]+ representing the operators star ({}), plus ({}), union ({}), intersection ({}), concatenation ({}) (default: {})'.format(''.join(OPERATOR_LIST), OPERATOR_STAR, OPERATOR_PLUS, OPERATOR_UNION, OPERATOR_INTER, OPERATOR_CONCAT, DEFAULT_OPERATORS)
    )
    regex_parser.add_argument(
        '--operator-type',
        '-O',
        dest    = 'operator_type',
        type    = str,
        choices = OPERATOR_TYPES,
        default = DEFAULT_OPERATOR_TYPE,
        help    = 'way to choose operator for deep terms (default: {})'.format(DEFAULT_OPERATOR_TYPE)
    )

    # random text fuzzer
    random_parser = subparsers.add_parser(RANDOM_TEXT, help='totally random text')
    random_parser.add_argument(
        '--length',
        '-l',
        dest    = 'length',
        metavar = 'L',
        type    = int,
        default = DEFAULT_LENGTH,
        help    = 'length of the text (default: {})'.format(DEFAULT_LENGTH)
    )

    # random_ast fuzzer
    random_ast_parser = subparsers.add_parser(RANDOM_AST, help='random but syntactically valid problem')
    random_ast_parser.add_argument(
        '--num-vars',
        '-v',
        dest    = 'num_vars',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_RANDOM_VARS,
        help    = 'number of variables (of each sort) in the problem (default: {})'.format(DEFAULT_NUM_RANDOM_VARS)
    )
    random_ast_parser.add_argument(
        '--num-asserts',
        '-n',
        dest    = 'num_asserts',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_RANDOM_ASSERTS,
        help    = 'number of asserts in the problem (default: {})'.format(DEFAULT_NUM_RANDOM_ASSERTS)
    )
    random_ast_parser.add_argument(
        '--depth',
        '-d',
        dest    = 'depth',
        metavar = 'D',
        type    = int,
        default = DEFAULT_RANDOM_DEPTH,
        help    = 'depth of nested expressions (default: {})'.format(DEFAULT_RANDOM_DEPTH)
    )
    random_ast_parser.add_argument(
        '--max-terms',
        '-t',
        dest    = 'max_terms',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MAX_RANDOM_TERMS,
        help    = 'maximum number of terms for n-ary expressions, like concats (default: {})'.format(DEFAULT_MAX_RANDOM_TERMS)
    )
    random_ast_parser.add_argument(
        '--max-string',
        '-l',
        dest    = 'max_str_lit_length',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MAX_RANDOM_STRINGS,
        help    = 'maximum length of string literals (default: {})'.format(DEFAULT_MAX_RANDOM_STRINGS)
    )
    random_ast_parser.add_argument(
        '--max-int',
        '-x',
        dest    = 'max_int_lit',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MAX_RANDOM_NUMBERS,
        help    = 'maximum size of int literals (default: {})'.format(DEFAULT_MAX_RANDOM_NUMBERS)
    )
    random_ast_parser.add_argument(
        '--meaningful',
        '-m',
        dest    = 'semantically_valid',
        action  = 'store_true',
        default = DEFAULT_SEMANTICALLY_VALID,
        help    = 'generate semantically valid problems (default: {})'.format(DEFAULT_SEMANTICALLY_VALID)
    )
    random_ast_parser.add_argument(
        '--literal-chance',
        '-p',
        dest    = 'literal_probability',
        metavar = 'P',
        type    = float,
        default = DEFAULT_LITERAL_PROBABILITY,
        help    = 'probability of creating literals instead of variables (default: {})'.format(DEFAULT_LITERAL_PROBABILITY)
    )

    # parse args
    args = global_parser.parse_args()

    # get the generator function based on args
    generator_name = args.generator
    generator      = GENERATORS[generator_name]

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

    # get some flags that will get popped from args before they're used
    produce_models = args.produce_models
    language       = args.language

    # get args as a dict
    # NOTE:
    #      argparse's Namespace object (which 'args' is) returns itself as a
    #      dict when vars() is called on it
    generator_args = vars(args)

    # pop arguments that are specific to this script because
    # they shouldn't be passed on to the generator
    generator_args.pop('language')
    generator_args.pop('produce_models')
    generator_args.pop('generator')
    generator_args.pop('seed')
    generator_args.pop('random')

    # run the generator with the args
    generated = generator(**generator_args)

    # prepend the logic setting
    generated = [smt_string_logic()] + generated

    # the random text generator produces raw text
    if (generator == random_text):
        print(generated)

    # other generators produce ASTs
    else:

        # add the model-getting node if needed
        if produce_models is True:
            generated.append(smt_get_model())

        print(generate(generated, language))

if __name__ == '__main__':
    main()
