import os
from random import randrange, choice, random
from typing import Callable

import tqdm
from frozendict import frozendict

from instrument import load_benchmark, Arg, Params, functions, invoke, call_statement, BranchTransformer, module_of

Range = tuple[int, int]

INT_RANGE: Range = (-1000, 1000)
STRING_LEN_RANGE: Range = (0, 10)
STRING_CHAR_RANGE: Range = (32, 127)
POOL_SIZE: int = 1000

OUT_DIR = os.path.join(os.path.dirname(__file__), "tests")


def random_int() -> int:
    return randrange(INT_RANGE[0], INT_RANGE[1])


def random_chr() -> str:
    chr_from, chr_to = STRING_CHAR_RANGE
    return chr(randrange(chr_from, chr_to))


def random_str() -> str:
    length = randrange(STRING_LEN_RANGE[0], STRING_LEN_RANGE[1])
    return "".join([random_chr() for _ in range(length)])


def max_cases(args: list[Arg]) -> int:
    num = 1
    for _, arg_type in args:
        if arg_type == 'int':
            num *= (INT_RANGE[1] - INT_RANGE[0])
        elif arg_type == 'str':
            len_from, len_to = STRING_LEN_RANGE
            chr_from, chr_to = STRING_CHAR_RANGE
            num *= sum([(chr_to - chr_from) * length * length for length in range(len_from, len_to)])
        else:
            raise ValueError(f"Arg type '{arg_type}' not supported")
    return num


def random_arg(arg_type: str) -> any:
    if arg_type == 'str':
        return random_str()
    elif arg_type == 'int':
        return random_int()
    else:
        raise ValueError(f"Arg type '{arg_type}' not supported")


def random_mutate(arg_type: str, arg_value: any) -> any:
    if arg_type == 'str':
        if len(arg_value) == 0:
            return arg_value

        prob = 1.0 / len(arg_value)
        for pos in range(len(arg_value)):
            if random() < prob:
                arg_value = list(arg_value)
                arg_value[pos] = random_chr()
                arg_value = "".join(arg_value)

        return arg_value
    elif arg_type == 'int':
        delta = randrange(-10, 10)
        return arg_value + delta
    else:
        raise ValueError(f"Arg type '{arg_type}' not supported")


def random_params(arguments: list[Arg]) -> Params:
    test_input: dict[str, any] = {}

    for arg_name, arg_type in arguments:
        test_input[arg_name] = random_arg(arg_type)

    return frozendict(test_input)


pools: dict[tuple, set[tuple]] = {}


def get_pool(arguments: list[Arg]) -> set[Params]:
    arg_types = tuple([arg_type for _, arg_type in arguments])
    arg_names = [arg_name for arg_name, _ in arguments]

    # Generate pool if not generated already
    # The pool only remembers the order of parameters and not their names
    if arg_types not in pools:
        new_pool = set()
        for _ in range(POOL_SIZE):
            param_list: list[any] = [None] * len(arg_names)

            params = random_params(arguments)
            for i, name in enumerate(arg_names):
                param_list[i] = params[name]

            new_pool.add(tuple(param_list))

        pools[arg_types] = new_pool

    return set([frozendict({arg_names[i]: p for i, p in enumerate(param)}) for param in pools[arg_types]])


def mutate(test_case: Params, arguments: list[Arg]) -> Params:
    arg_name = choice(list(test_case.keys()))  # choose name to mutate
    types: dict[str, str] = {arg_name: arg_type for arg_name, arg_type in arguments}
    return test_case.set(arg_name, random_mutate(types[arg_name], test_case[arg_name]))


def crossover(chosen_test: Params, other_chosen_test: Params, arguments: list[Arg]) -> tuple[Params, Params]:
    # Select a property at random and swap properties
    arg_name = choice(list(chosen_test.keys()))
    types: dict[str, str] = {arg_name: arg_type for arg_name, arg_type in arguments}
    if types[arg_name] == 'str':
        # Crossover for strings intermingles the strings of the two chosen tests
        s1, s2 = str_crossover(chosen_test[arg_name], other_chosen_test[arg_name])
        t1 = chosen_test.set(arg_name, s1)
        t2 = other_chosen_test.set(arg_name, s2)

    else:  # types[arg_name] == 'int'
        # Crossover for integers swaps the values from the two tests
        i1, i2 = chosen_test[arg_name], other_chosen_test[arg_name]
        t1 = chosen_test.set(arg_name, i1)
        t2 = other_chosen_test.set(arg_name, i2)

    return t1, t2


def get_test_cases(f_name: str, arguments: list[Arg], n: int, enable_bar=True) -> set[Params]:
    assert n >= 1

    pool: set[Params] = get_pool(arguments)
    pool_list = list(pool)
    tests: set[Params] = set()

    n = min(n, max_cases(arguments) // 3)  # bound n by 1/3rd of the max possible number of tests

    with tqdm.tqdm(total=n, desc=f"Tests for {BranchTransformer.to_original_name(f_name)}",
                   disable=not enable_bar) as pbar:
        def consider_test_case(t: Params):
            if t not in pool:
                pool.add(t)
                pool_list.append(t)

            try:
                invoke(f_name, t)  # check if this input satisfies the input assertion
            except AssertionError:
                return

            if t not in tests:
                tests.add(t)
                pbar.update()

        while len(tests) < n:
            chosen_test: Params = choice(pool_list)
            kind = choice(['pool', 'mutation', 'crossover'])

            if kind == 'mutation':
                consider_test_case(mutate(chosen_test, arguments))
            elif kind == 'crossover':
                # pick other distinct sample
                while True:
                    other_chosen_test: Params = choice(pool_list)
                    if frozendict(chosen_test) != frozendict(other_chosen_test):
                        break

                t1, t2 = crossover(chosen_test, other_chosen_test, arguments)
                consider_test_case(t1)
                consider_test_case(t2)
            else:
                consider_test_case(chosen_test)

    return tests


def str_crossover(parent1: str, parent2: str):
    if len(parent1) > 1 and len(parent2) > 1:
        pos = randrange(1, len(parent1))
        offspring1 = parent1[:pos] + parent2[pos:]
        offspring2 = parent2[:pos] + parent1[pos:]
        return offspring1, offspring2

    return parent1, parent2


def get_test_case_source(f_name: str, test_case: Params, i: int, indent: int):
    f_name_orig = BranchTransformer.to_original_name(f_name)

    single_indent = " " * 4
    space = single_indent * indent

    output = invoke(f_name, test_case)

    if type(output) == str:
        output = f"'{output}'"

    return f"""{space}def test_{f_name_orig}_{i}(self):
{space}{single_indent}assert {call_statement(f_name_orig, test_case)} == {output}"""


def get_test_class(f_name: str, case: set[Params]) -> str:
    f_name_orig = BranchTransformer.to_original_name(f_name)

    test_class = (f"from unittest import TestCase\n\nfrom {module_of[f_name]} import {f_name_orig}\n\n\n"
                  f"class Test_{f_name_orig}(TestCase):\n")
    test_class += "\n\n".join([get_test_case_source(f_name, case, i + 1, 1) for i, case in enumerate(cases)])
    return test_class


def main():
    load_benchmark(save_instrumented=False)  # instrument all files in benchmark

    if not os.path.isdir(OUT_DIR):
        os.makedirs(OUT_DIR)

    for f_name in functions.keys():
        with open(os.path.join(OUT_DIR, f_name + ".py"), "w") as f:
            cases = get_test_cases(f_name, functions[f_name], 100)
            f.write(get_test_class(f_name, cases))


if __name__ == '__main__':
    main()