# __main__.py
import logging
import sys
from anubis import arg_parser_main, results
from anubis.feature_file_parser import get_tests
from anubis.parallelizer import test_runner
from anubis.text import print_console_output
from behave.model import Feature, Scenario, ScenarioOutline
from datetime import datetime
from multiprocessing import Pool
from types import SimpleNamespace


def main() -> int:
    """Orchestrate parallel execution of BDD tests."""
    # Setup -----------------------------------------------------------------------------------------------------------
    start: datetime = datetime.now()
    args, args_unknown = arg_parser_main.parse_arguments()
    print_console_output('startup_statement', **vars(args))

    # Set up logging
    args.output.mkdir(parents=True, exist_ok=True)
    args.log_file.parent.mkdir(parents=True, exist_ok=True)

    log_handlers = [logging.FileHandler(args.log_file, mode="w+")]
    if args.log_std_out:
        log_handlers.append(logging.StreamHandler(sys.stdout))
    logging.basicConfig(
        handlers=log_handlers,
        level=args.log_level,
        format='[%(name)s] [%(asctime)s] [%(levelname)s] %(message)s',
        datefmt='%H:%M:%S'
    )
    logger: logging.Logger = logging.getLogger()
    logger.info('Args for this test run:\n\t' + "\n\t".join(str(a) for a in vars(args).items()))

    # Parse and prepare tests
    print_console_output('output_statement', **vars(args))
    print_console_output('parameter_statement', **vars(args))
    logger.info('parsing tests')

    tests_to_run: list[Feature | Scenario | ScenarioOutline] = get_tests(
        args.features, args.tags, args.unit
    )

    # Handle dry-run mode
    if args.dry_run:
        test_split = _distribute_tests(tests_to_run, args.processes)
        print_console_output('dry_run_statement', test_split=test_split, **vars(args))
        return 0

    # Execute tests -------------------------------------------------------------------------------------------------
    print_console_output('running_statement', test_split={}, **vars(args))
    logger.info('starting test pool')

    test_split = _distribute_tests(tests_to_run, args.processes)
    output_files = _run_tests_parallel(test_split, args, args_unknown)

    # Process results ------------------------------------------------------------------------------------------------
    results.write_result_aggregate(files=output_files, aggregate_out_file=args.aggregate)
    results.json_results_to_xml(start, args.aggregate, args.junit)
    logger.info(f'output files: {output_files}')

    passed, failed, total = results.get_result_values(args.aggregate)
    results.print_result_summary(args, start, datetime.now(), passed, failed)
    logger.info(f'passed: {passed}, failed: {failed}')

    # Cleanup and exit ------------------------------------------------------------------------------------------------
    if args.delete_output:
        args.output.rm_dir()

    print_console_output('end_statement', **vars(args))
    return _get_exit_code(passed, total, args.pass_threshold, args.pass_with_no_tests)


def _distribute_tests(tests: list, num_processes: int) -> dict:
    """Distribute tests evenly across worker processes using round-robin."""
    test_split: dict = {i: [] for i in range(num_processes)}
    for i, test in enumerate(tests):
        test_split[i % num_processes].append(test)
    return test_split


def _run_tests_parallel(test_split: dict, args, args_unknown: list) -> list:
    """Execute tests in parallel using multiprocessing pool."""
    with Pool(processes=len(test_split)) as pool:
        run_args = [
            (
                i,
                SimpleNamespace(output=args.output, D=args.D, tags=args.tags),
                args_unknown,
                [t.location.filename for t in tests],
            )
            for i, tests in test_split.items()
        ]
        output_files = pool.starmap(test_runner, run_args)
    return output_files


def _get_exit_code(passed: int, total: int, pass_threshold: float, pass_with_no_tests: bool) -> int:
    """Calculate exit code based on test results.
    
    Returns:
        0: All tests passed (or no tests with pass_with_no_tests=True)
        1: Tests failed or below pass_threshold
    """
    if total == 0:
        return 0 if pass_with_no_tests else 1
    return 0 if (passed / total >= pass_threshold) else 1


if __name__ == '__main__':
    sys.exit(main())
