from itertools import product

from shrub.v3.evg_command import KeyValueParam, expansions_update
from shrub.v3.evg_task import EvgTask, EvgTaskDependency

from config_generator.components.funcs.bootstrap_mongo_orchestration import BootstrapMongoOrchestration
from config_generator.components.funcs.csfle_setup import CSFLESetup
from config_generator.components.funcs.csfle_teardown import CSFLETeardown
from config_generator.components.funcs.fetch_build import FetchBuild
from config_generator.components.funcs.fetch_det import FetchDET
from config_generator.components.funcs.run_tests import RunTests
from config_generator.etc.distros import compiler_to_vars, find_large_distro, find_small_distro, make_distro_str


def generate_test_tasks(SSL, TAG, MATRIX):
    res = []

    TOPOLOGY_TO_STR = {
        'server': 'server',
        'replica': 'replica_set',
        'sharded': 'sharded_cluster',
    }

    for distro_name, compiler, arch, sasl, auths, topologies, server_vers in MATRIX:
        tags = [TAG, 'test', distro_name, compiler, f'sasl-{sasl}', 'cse']
        if distro_name == 'rhel8-latest':
            test_distro = find_large_distro(distro_name)  # DEVPROD-18763
        else:
            test_distro = find_small_distro(distro_name)

        compile_vars = []
        compile_vars = [KeyValueParam(key=key, value=value) for key, value in compiler_to_vars(compiler).items()]

        if arch:
            tags.append(arch)
            compile_vars.append(KeyValueParam(key='MARCH', value=arch))

        distro_str = make_distro_str(distro_name, compiler, arch)
        compile_task_name = f'cse-sasl-{sasl}-{SSL}-{distro_str}-compile'

        for auth, topology, server_ver in product(auths, topologies, server_vers):
            test_tags = tags + [auth, topology, server_ver]
            # Do not add `nossl` tag to prevent being selected by legacy config variants.
            # Remove the `if` when CDRIVER-4571 is resolved.
            if SSL != 'nossl':
                test_tags += [SSL]
            test_task_name = f'cse-sasl-{sasl}-{SSL}-{distro_str}-test-{server_ver}-{topology}-{auth}'

            test_commands = []
            test_commands.append(FetchBuild.call(build_name=compile_task_name))

            updates = compile_vars + [
                KeyValueParam(key='AUTH', value=auth),
                KeyValueParam(key='MONGODB_VERSION', value=server_ver),
                KeyValueParam(key='TOPOLOGY', value=TOPOLOGY_TO_STR[topology]),
                KeyValueParam(key='SSL', value=SSL),
            ]

            test_commands.append(expansions_update(updates=updates))
            test_commands.append(FetchDET.call())
            test_commands.append(BootstrapMongoOrchestration.call())
            test_commands.append(CSFLESetup.call())
            test_commands.append(RunTests.call())
            test_commands.append(CSFLETeardown.call())

            res.append(
                EvgTask(
                    name=test_task_name,
                    run_on=test_distro.name,
                    tags=test_tags,
                    depends_on=[EvgTaskDependency(name=compile_task_name)],
                    commands=test_commands,
                )
            )

    return res
