#!/usr/bin/env python3

'''
Merge two existing problems.
'''

import sys
import argparse
import random

from stringfuzz.constants import LANGUAGES, SMT_25_STRING
from stringfuzz.generator import generate
from stringfuzz.parser import parse, ParsingError
from stringfuzz.smt import smt_string_logic, smt_check_sat
from stringfuzz.ast import SettingNode, ExpressionNode, MetaCommandNode, GenericExpressionNode, CheckSatNode, GetModelNode
from stringfuzz.mergers import simple, DISJOINT_RENAME, INTERSECTING_RENAME, NO_RENAME

#constants
SIMPLE = 'simple'

MERGERS = {
    SIMPLE: simple
}

RENAMING_STRATEGIES = [
    DISJOINT_RENAME,
    INTERSECTING_RENAME,
    NO_RENAME,
]

# defaults
DEFAULT_RENAMING = NO_RENAME
DEFAULT_SEED     = 0
DEFAULT_RANDOM   = False

GET_MODEL     = "get-model"
GET_INFO      = "get-info"
EXPR_TO_STRIP = [GET_MODEL, GET_INFO]
CHECK_SAT     = "check-sat"
GEN_TO_STRIP  = [CHECK_SAT]

def should_keep(node):
    if isinstance(node, SettingNode):
        return False
    if isinstance(node, MetaCommandNode):
        return False
    if isinstance(node, CheckSatNode):
        return False
    if isinstance(node, GetModelNode):
        return False
    if isinstance(node, ExpressionNode):
        if node.symbol in EXPR_TO_STRIP:
            return False
    if isinstance(node, GenericExpressionNode):
        if node.symbol.name in GEN_TO_STRIP:
            return False
    return True

# entry point
def main():

    # create arg parser
    global_parser = argparse.ArgumentParser(description='SMTLIB 2.* problem merger.')
    global_parser.add_argument(
        'files',
        nargs    = '+',
        metavar  = 'F',
        type     = argparse.FileType('r'),
        help     = 'input files'
    )
    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='merger', help='merger choice')
    subparsers.required = True

    # simple merger
    simple_parser = subparsers.add_parser(SIMPLE, help='simple merger')
    simple_parser.add_argument(
        '--renaming',
        dest    = 'rename_type',
        type    = str,
        choices = RENAMING_STRATEGIES,
        default = DEFAULT_RENAMING,
        help    = 'variable renaming strategy (default: {})'.format(DEFAULT_RENAMING)
    )

    # parse args
    args = global_parser.parse_args()

    # get the merger function based on args
    merger_name = args.merger
    merger      = MERGERS[merger_name]

    # get general args
    files           = args.files
    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)

    # get args as a dict and pop general args
    merge_args = vars(args)
    merge_args.pop('files')
    merge_args.pop('input_language')
    merge_args.pop('output_language')
    merge_args.pop('seed')
    merge_args.pop('random')
    merge_args.pop('merger')

    # read input
    raw_in = [f.read() for f in files]

    # parse input
    try:
        asts = [parse(raw, input_language) for raw in raw_in]
    except ParsingError as e:
        print(e, file=sys.stderr)
        return 1

    # filter out suppressed expressions
    asts = [list(filter(should_keep, ast)) for ast in asts]

    # merge the two ASTs into a new AST
    merged = merger(asts, **merge_args)

    # add back the logic and get-sat
    merged = [smt_string_logic()] + merged + [smt_check_sat()]

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

if __name__ == '__main__':
    main()
