#!/usr/bin/env python3
# Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
#
# See LICENSE.txt for license information
"""
Test Parser Module
Handles command-line argument parsing and test output parsing
"""

import re
import argparse


class ArgumentParserInterface:
    """Command-line argument parser for RCCL test runner"""

    def __init__(self):
        self.parser = argparse.ArgumentParser(
            description="RCCL Test Runner - Execute and manage RCCL unit tests and MPI tests",
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog="""
Examples:
  # Run all tests from config
  %(prog)s -c test_config.json

  # Run a specific individual test by exact name
  %(prog)s -c test_config.json --test-name CE_AllGather_2Ranks

  # Run multiple individual tests (colon-separated)
  %(prog)s -c test_config.json --test-name CE_AllGather_2Ranks:CE_AlltoAll_2Ranks

  # Run all suites whose name contains 'CE Tests'
  %(prog)s -c test_config.json --suite-name "CE Tests"

  # Run two specific suites by substring (colon-separated)
  %(prog)s -c test_config.json --suite-name "CE Tests - 2-Rank:CE Tests - 4-Rank"

  # Run with verbose output
  %(prog)s -c test_config.json -v

  # Skip build and use existing build
  %(prog)s -c test_config.json --no-build

  # Use a custom build directory
  %(prog)s -c test_config.json --build-dir /tmp/my_rccl_build

  # Generate coverage report from existing data
  %(prog)s -c test_config.json --no-build --skip-tests --coverage-report

  # Run on thor2 system (uses thor2 MPI args from config)
  %(prog)s -c test_config.json --no-build --system thor2
            """
        )

    def add_arguments(self):
        """Add all command-line arguments"""
        self.parser.add_argument(
            '-c', '--config',
            type=str,
            required=True,
            help="Test configuration file (JSON format)"
        )
        self.parser.add_argument(
            '-v', '--verbose',
            action='store_true',
            help="Enable verbose output (detailed logging)"
        )
        self.parser.add_argument(
            '-o', '--output',
            type=str,
            help="Output directory for logs and reports (default: auto-generated)"
        )
        self.parser.add_argument(
            '--test-name',
            type=str,
            help="Run only tests matching any glob pattern; ':' = OR, '*' = wildcard, '-' prefix = exclude, case-sensitive (e.g. 'P2P_*' = all P2P tests; '*:-SHM*' = all except SHM)"
        )
        self.parser.add_argument(
            '--suite-name',
            type=str,
            help="gtest-style glob filter on suite names: ':' = OR, '*' = wildcard, '-' prefix = exclude, case-sensitive (e.g. 'CE*' runs all CE suites; '*:-NET*' runs all except NET)"
        )
        self.parser.add_argument(
            '--no-build',
            action='store_true',
            help="Skip build step and use existing build artifacts"
        )
        self.parser.add_argument(
            '--skip-tests',
            action='store_true',
            help="Skip test execution (useful with --coverage-report)"
        )
        self.parser.add_argument(
            '--coverage-report',
            action='store_true',
            help="Generate code coverage report from profraw files"
        )
        self.parser.add_argument(
            '--build-dir',
            type=str,
            help="Custom build directory path (default: <workdir>/build/debug or build/release)"
        )
        self.parser.add_argument(
            '--report-suffix',
            type=str,
            default='',
            help="Suffix for report directory name (default: blank)"
        )
        self.parser.add_argument(
            '--rerun-failed',
            action='store_true',
            help="Rerun failed tests with additional environment variables from config (rerun_env_variables)"
        )
        self.parser.add_argument(
            '--skip-mpi-check',
            action='store_true',
            help="Skip MPI: removes --enable-mpi-tests from build flags, skips MPI installation check, and skips tests with num_ranks > 1"
        )
        self.parser.add_argument(
            '--stop-on-rerun-failure',
            action='store_true',
            help="Stop testing immediately if a rerun also fails (requires --rerun-failed)"
        )
        self.parser.add_argument(
            '--system',
            type=str,
            default='',
            help="Select system-specific MPI args profile from config (e.g. 'ainic', 'thor2')"
        )
        self.parser.add_argument(
            '--mpich',
            action='store_true',
            default=False,
            help="Use MPICH syntax (-env) instead of OpenMPI (-x) for passing env vars to mpirun"
        )

    def parse_arguments(self):
        """Parse command-line arguments"""
        return self.parser.parse_args()

    def process_arguments(self):
        """Process and validate command-line arguments"""
        self.add_arguments()
        args = self.parse_arguments()
        self.handle_arguments(args)
        return args

    def handle_arguments(self, args):
        """Handle and display parsed arguments"""
        if args.verbose:
            print("="*80)
            print("RCCL Test Runner - Configuration")
            print("="*80)
            print(f"Config file:       {args.config}")
            print(f"Verbose mode:      {args.verbose}")
            print(f"Output dir:        {args.output if args.output else 'auto-generated'}")
            print(f"Test name filter:  {args.test_name if args.test_name else 'all (glob)'}")
            print(f"Suite name filter: {args.suite_name if args.suite_name else 'all suites'}")
            print(f"No build:          {args.no_build}")
            print(f"Skip tests:        {args.skip_tests}")
            print(f"Coverage report:   {args.coverage_report}")
            print(f"Build dir:         {args.build_dir if args.build_dir else 'default'}")
            print(f"Report suffix:     {args.report_suffix}")
            print(f"Rerun failed:      {args.rerun_failed}")
            print(f"Skip MPI check:    {args.skip_mpi_check}")
            print(f"Stop on rerun fail: {args.stop_on_rerun_failure}")
            print(f"System profile:    {args.system if args.system else 'none (use default MPI args)'}")
            print(f"MPI implementation: {'mpich' if args.mpich else 'openmpi (default)'}")
            print("="*80)
            print()


def parse_test_output(output):
    """
    Parse test output and extract results

    Args:
        output: String containing test output

    Returns:
        dict: Parsed test results including pass/fail status
    """
    results = {
        'passed': False,
        'failed': False,
        'skipped': False,
        'tests_run': 0,
        'tests_passed': 0,
        'tests_failed': 0,
        'errors': []
    }

    # Google Test output patterns
    gtest_passed = re.search(r'\[\s*PASSED\s*\]\s*(\d+)\s*test', output)
    gtest_failed = re.search(r'\[\s*FAILED\s*\]\s*(\d+)\s*test', output)
    gtest_run = re.search(r'\[==========\]\s*(\d+)\s*test.*ran', output)

    if gtest_run:
        results['tests_run'] = int(gtest_run.group(1))

    if gtest_passed:
        results['tests_passed'] = int(gtest_passed.group(1))

    if gtest_failed:
        results['tests_failed'] = int(gtest_failed.group(1))
        results['failed'] = True
    else:
        results['passed'] = results['tests_run'] > 0

    # Check for skipped tests
    if 'SKIPPED' in output or 'Skipped' in output:
        results['skipped'] = True

    # Extract error messages
    error_pattern = re.compile(r'(ERROR|FAILED|TIMEOUT).*', re.MULTILINE)
    errors = error_pattern.findall(output)
    results['errors'] = errors[:10]  # Limit to first 10 errors

    return results

