#!/usr/bin/env python3

import sys
import os
import argparse

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

from stringfuzz.fuzzers.genetic import simulate, SAT_ANSWER, UNSAT_ANSWER

# constants
DEFAULT_NUM_GENERATIONS = 200
DEFAULT_WORLD_SIZE      = 10
DEFAULT_TIMEOUT         = 5

AVAILABLE_WANTED_ANSWERS = [
    SAT_ANSWER,
    UNSAT_ANSWER
]

DEFAULT_MIN_ASSERTS      = 3
DEFAULT_MAX_ASSERTS      = 10
DEFAULT_WANTED_ANSWER    = None
DEFAULT_TRACE_RESOLUTION = 1
TRACE_BY_DEFAULT         = False

DEFAULT_SEED_PROBLEM = None

# globals
_tracing_on = False

# helpers
def trace(*args, **kwargs):
    global _tracing_on
    if _tracing_on is True:
        print(*args, file=sys.stderr, **kwargs)

# main
def main():

    # create arg parser
    parser = argparse.ArgumentParser(description='Test an SMT string solver in an exploratory fashion.')
    parser.add_argument(
        'command',
        type = str,
        help = 'command to run the solver (if input fed on standard in)'
    )

    # I/O arguments
    parser.add_argument(
        '--out-language',
        '-o',
        dest    = 'out_language',
        type    = str,
        choices = LANGUAGES,
        default = SMT_25_STRING,
        help    = 'solver language (default: {})'.format(SMT_25_STRING)
    )
    parser.add_argument(
        '--in-language',
        '-i',
        dest    = 'in_language',
        type    = str,
        choices = LANGUAGES,
        default = SMT_25_STRING,
        help    = 'seed problem language (default: {})'.format(SMT_25_STRING)
    )
    parser.add_argument(
        '--trace',
        '-v',
        dest    = 'trace',
        action  = 'store_true',
        default = TRACE_BY_DEFAULT,
        help    = 'print status messages (default: {})'.format(TRACE_BY_DEFAULT)
    )
    parser.add_argument(
        '--seed-problem',
        '-s',
        dest    = 'seed_problem',
        metavar = 'F',
        default = DEFAULT_SEED_PROBLEM,
        type    = argparse.FileType('r'),
        help    = 'input file (default: {})'.format(DEFAULT_SEED_PROBLEM)
    )

    # simulation arguments
    parser.add_argument(
        '--num-generations',
        '-g',
        dest    = 'num_generations',
        metavar = 'N',
        type    = int,
        default = DEFAULT_NUM_GENERATIONS,
        help    = 'number of generations to simulate (default: {})'.format(DEFAULT_NUM_GENERATIONS)
    )
    parser.add_argument(
        '--world-size',
        '-w',
        dest    = 'world_size',
        metavar = 'N',
        type    = int,
        default = DEFAULT_WORLD_SIZE,
        help    = 'number of organisms in the world (default: {})'.format(DEFAULT_WORLD_SIZE)
    )
    parser.add_argument(
        '--min-asserts',
        '-m',
        dest    = 'min_asserts',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MIN_ASSERTS,
        help    = 'min number of asserts per instance (except the seed) (default: {})'.format(DEFAULT_MIN_ASSERTS)
    )
    parser.add_argument(
        '--max-asserts',
        '-M',
        dest    = 'max_asserts',
        metavar = 'N',
        type    = int,
        default = DEFAULT_MAX_ASSERTS,
        help    = 'max number of asserts per instance (except the seed) (default: {})'.format(DEFAULT_MAX_ASSERTS)
    )
    parser.add_argument(
        '--timeout',
        '-t',
        dest    = 'timeout',
        type    = float,
        default = DEFAULT_TIMEOUT,
        help    = 'timeout, in seconds (default: {})'.format(DEFAULT_TIMEOUT)
    )
    parser.add_argument(
        '--answer',
        '-a',
        dest    = 'wanted_answer',
        type    = str,
        choices = AVAILABLE_WANTED_ANSWERS,
        default = DEFAULT_WANTED_ANSWER,
        help    = 'the answer which to prefer (default: {})'.format(DEFAULT_WANTED_ANSWER)
    )

    # parse args
    args = parser.parse_args()

    # set tracing mode
    global _tracing_on
    _tracing_on = args.trace

    # use passed seed problem
    if args.seed_problem is not None:
        seed_problem = parse(args.seed_problem.read(), args.in_language)

    # or create a random one
    else:

        # figure out how many asserts to generate
        half_max_asserts   = args.max_asserts // 2
        num_random_asserts = max(1, half_max_asserts)

        seed_problem = [smt_string_logic()] + random_ast(
            num_vars            = 1,
            num_asserts         = num_random_asserts,
            depth               = 3,
            max_terms           = 5,
            max_str_lit_length  = 10,
            max_int_lit         = 30,
            literal_probability = 0.5,
            semantically_valid  = True
        )

    # print seed problem
    trace('-----')
    trace('solver command:', args.command)
    trace('-----')
    trace('progenitor:')
    trace('-----')
    trace(generate(seed_problem, args.out_language))
    trace('-----')

    # run the tester
    trace('')
    trace('simulating ...')
    surviving_instances = simulate(
        progenitor       = seed_problem,
        language         = args.out_language,
        solver_command   = args.command,
        num_generations  = args.num_generations,
        world_size       = args.world_size,
        solver_timeout   = args.timeout,
        wanted_answer    = args.wanted_answer,
        max_num_asserts  = args.max_asserts,
        min_num_asserts  = args.min_asserts,
        trace_resolution = DEFAULT_TRACE_RESOLUTION,
        enable_tracing   = _tracing_on
    )
    trace('finished')
    trace('')

    # print last survivor
    renegade      = surviving_instances[0]
    renegade_text = generate(renegade, args.out_language)
    print(renegade_text)

if __name__ == '__main__':
    main()
