#!/usr/bin/env python3
"""
Unified Test Runner for PrecipGen Workflow Testing Framework

This module provides a comprehensive test execution framework that runs all workflow tests
with organized reporting, visualization, and automated execution capabilities.
"""

import os
import sys
import json
import time
import subprocess
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field, asdict
from enum import Enum
import argparse

# Add project root to path for imports
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from tests.workflow_testing.infrastructure import WorkflowTestLogger, TestConfig, DirectoryManager


class TestStatus(Enum):
    """Test execution status."""
    NOT_STARTED = "not_started"
    RUNNING = "running"
    PASSED = "passed"
    FAILED = "failed"
    SKIPPED = "skipped"
    ERROR = "error"


@dataclass
class TestResult:
    """Result from a single test execution."""
    test_name: str
    test_file: str
    status: TestStatus
    execution_time: float = 0.0
    error_message: Optional[str] = None
    output: Optional[str] = None
    property_test_results: Dict[str, Any] = field(default_factory=dict)
    performance_metrics: Dict[str, float] = field(default_factory=dict)


@dataclass
class WorkflowTestSuite:
    """Results from a complete workflow test suite."""
    workflow_name: str
    test_results: List[TestResult] = field(default_factory=list)
    total_tests: int = 0
    passed_tests: int = 0
    failed_tests: int = 0
    skipped_tests: int = 0
    error_tests: int = 0
    total_execution_time: float = 0.0
    success_rate: float = 0.0


@dataclass
class TestExecutionReport:
    """Complete test execution report."""
    execution_id: str
    start_time: datetime
    end_time: Optional[datetime] = None
    total_execution_time: float = 0.0
    workflow_suites: List[WorkflowTestSuite] = field(default_factory=list)
    overall_success_rate: float = 0.0
    total_tests: int = 0
    total_passed: int = 0
    total_failed: int = 0
    total_skipped: int = 0
    total_errors: int = 0
    configuration: Dict[str, Any] = field(default_factory=dict)
    environment_info: Dict[str, str] = field(default_factory=dict)


class WorkflowTestRunner:
    """Unified test runner for all PrecipGen workflow tests."""
    
    def __init__(self, output_dir: Optional[str] = None, config: Optional[TestConfig] = None):
        """Initialize the test runner.
        
        Args:
            output_dir: Directory for test outputs and reports
            config: Test configuration (uses default if None)
        """
        self.config = config or TestConfig()
        self.output_dir = Path(output_dir or "./test_results")
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # Initialize directory manager for organized outputs
        self.dir_manager = DirectoryManager(base_dir=str(self.output_dir))
        
        # Test discovery
        self.test_directory = Path(__file__).parent
        self.workflow_tests = self._discover_workflow_tests()
        
        # Execution tracking
        self.execution_id = f"test_run_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.logger = None
        
    def _discover_workflow_tests(self) -> Dict[str, List[str]]:
        """Discover all workflow test files organized by workflow type."""
        workflow_tests = {
            "directory_organization": ["test_directory_organization.py"],
            "data_acquisition": ["test_data_acquisition.py", "test_data_download.py", "test_quality_assessment.py"],
            "data_analysis": ["test_data_analysis.py"],
            "simulation": ["test_simulation_workflow.py"],
            "validation": ["test_validation_workflow.py"],
            "integration": ["test_integration_workflow.py"],
            "performance": ["test_performance_simple.py", "test_performance_scalability.py", "test_batch_processing.py"],
            "error_handling": ["test_error_handling.py", "test_error_recovery_property.py"],
            "output_format": ["test_output_format_standardization.py"]
        }
        
        # Filter to only include existing test files
        existing_tests = {}
        for workflow, test_files in workflow_tests.items():
            existing_files = []
            for test_file in test_files:
                test_path = self.test_directory / test_file
                if test_path.exists():
                    existing_files.append(test_file)
            if existing_files:
                existing_tests[workflow] = existing_files
                
        return existing_tests
    
    def _get_environment_info(self) -> Dict[str, str]:
        """Collect environment information for the report."""
        import platform
        import precipgen
        
        return {
            "python_version": platform.python_version(),
            "platform": platform.platform(),
            "precipgen_version": getattr(precipgen, "__version__", "unknown"),
            "working_directory": str(Path.cwd()),
            "test_directory": str(self.test_directory),
            "output_directory": str(self.output_dir)
        }
    
    def _run_single_test(self, test_file: str, workflow: str) -> TestResult:
        """Run a single test file and collect results."""
        test_path = self.test_directory / test_file
        test_name = test_file.replace(".py", "")
        
        result = TestResult(
            test_name=test_name,
            test_file=test_file,
            status=TestStatus.NOT_STARTED
        )
        
        try:
            result.status = TestStatus.RUNNING
            start_time = time.time()
            
            # Run pytest on the specific test file
            cmd = [
                sys.executable, "-m", "pytest", 
                str(test_path),
                "-v", "--tb=short",
                f"--junitxml={self.output_dir}/junit_{test_name}.xml"
            ]
            
            # Add property test configuration if it's a property test
            if "property" in test_name:
                cmd.extend(["--hypothesis-max-examples=100", "--hypothesis-deadline=10000"])
            
            process = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=self.config.get_workflow_timeout(workflow)
            )
            
            result.execution_time = time.time() - start_time
            result.output = process.stdout + process.stderr
            
            if process.returncode == 0:
                result.status = TestStatus.PASSED
            else:
                result.status = TestStatus.FAILED
                result.error_message = f"Test failed with return code {process.returncode}"
                
        except subprocess.TimeoutExpired:
            result.status = TestStatus.ERROR
            result.error_message = f"Test timed out after {self.config.get_workflow_timeout(workflow)} seconds"
            result.execution_time = self.config.get_workflow_timeout(workflow)
            
        except Exception as e:
            result.status = TestStatus.ERROR
            result.error_message = f"Test execution error: {str(e)}"
            result.execution_time = time.time() - start_time if 'start_time' in locals() else 0.0
            
        return result
    
    def _run_workflow_suite(self, workflow: str, test_files: List[str]) -> WorkflowTestSuite:
        """Run all tests for a specific workflow."""
        suite = WorkflowTestSuite(workflow_name=workflow)
        
        self.logger.logger.info(f"Starting workflow test suite: {workflow}")
        suite_start_time = time.time()
        
        for test_file in test_files:
            self.logger.logger.info(f"Running test: {test_file}")
            result = self._run_single_test(test_file, workflow)
            suite.test_results.append(result)
            
            # Update counters
            if result.status == TestStatus.PASSED:
                suite.passed_tests += 1
            elif result.status == TestStatus.FAILED:
                suite.failed_tests += 1
            elif result.status == TestStatus.SKIPPED:
                suite.skipped_tests += 1
            elif result.status == TestStatus.ERROR:
                suite.error_tests += 1
                
            suite.total_tests += 1
            
            # Log result
            self.logger.log_test_result(
                test_case=result.test_name,
                success=(result.status == TestStatus.PASSED),
                execution_time=result.execution_time,
                details={"status": result.status.value, "error": result.error_message}
            )
        
        suite.total_execution_time = time.time() - suite_start_time
        suite.success_rate = suite.passed_tests / suite.total_tests if suite.total_tests > 0 else 0.0
        
        self.logger.logger.info(f"Workflow {workflow} completed: {suite.passed_tests}/{suite.total_tests} passed ({suite.success_rate:.1%})")
        
        return suite
    
    def run_all_tests(self, workflows: Optional[List[str]] = None, 
                     parallel: bool = False) -> TestExecutionReport:
        """Run all workflow tests and generate comprehensive report.
        
        Args:
            workflows: Specific workflows to test (all if None)
            parallel: Whether to run workflows in parallel (not implemented yet)
            
        Returns:
            Complete test execution report
        """
        # Initialize execution report
        report = TestExecutionReport(
            execution_id=self.execution_id,
            start_time=datetime.now(),
            configuration=self.config.__dict__,
            environment_info=self._get_environment_info()
        )
        
        # Initialize logger
        log_dir = self.output_dir / "logs"
        log_dir.mkdir(exist_ok=True)
        
        with WorkflowTestLogger(self.execution_id, log_dir=str(log_dir)) as logger:
            self.logger = logger
            
            logger.start_test_execution(f"Running PrecipGen workflow tests: {self.execution_id}")
            
            # Determine which workflows to run
            workflows_to_run = workflows or list(self.workflow_tests.keys())
            
            # Run each workflow test suite
            for workflow in workflows_to_run:
                if workflow not in self.workflow_tests:
                    logger.logger.warning(f"Workflow '{workflow}' not found, skipping")
                    continue
                    
                test_files = self.workflow_tests[workflow]
                suite = self._run_workflow_suite(workflow, test_files)
                report.workflow_suites.append(suite)
                
                # Update overall counters
                report.total_tests += suite.total_tests
                report.total_passed += suite.passed_tests
                report.total_failed += suite.failed_tests
                report.total_skipped += suite.skipped_tests
                report.total_errors += suite.error_tests
            
            # Finalize report
            report.end_time = datetime.now()
            report.total_execution_time = (report.end_time - report.start_time).total_seconds()
            report.overall_success_rate = report.total_passed / report.total_tests if report.total_tests > 0 else 0.0
            
            logger.end_test_execution(
                success=(report.total_failed == 0 and report.total_errors == 0),
                summary=f"Completed {report.total_tests} tests: {report.total_passed} passed, {report.total_failed} failed, {report.total_errors} errors"
            )
        
        # Save detailed report
        self._save_report(report)
        
        return report
    
    def _save_report(self, report: TestExecutionReport) -> None:
        """Save the test execution report to files."""
        # Save JSON report
        json_report_path = self.output_dir / f"test_report_{self.execution_id}.json"
        with open(json_report_path, 'w') as f:
            json.dump(asdict(report), f, indent=2, default=str)
        
        # Save human-readable summary
        summary_path = self.output_dir / f"test_summary_{self.execution_id}.txt"
        with open(summary_path, 'w') as f:
            f.write(self._generate_summary_report(report))
        
        print(f"\nTest execution completed!")
        print(f"Detailed report: {json_report_path}")
        print(f"Summary report: {summary_path}")
    
    def _generate_summary_report(self, report: TestExecutionReport) -> str:
        """Generate a human-readable summary report."""
        lines = []
        lines.append("=" * 80)
        lines.append("PrecipGen Workflow Testing Framework - Execution Report")
        lines.append("=" * 80)
        lines.append(f"Execution ID: {report.execution_id}")
        lines.append(f"Start Time: {report.start_time}")
        lines.append(f"End Time: {report.end_time}")
        lines.append(f"Total Execution Time: {report.total_execution_time:.2f} seconds")
        lines.append("")
        
        # Overall summary
        lines.append("OVERALL RESULTS")
        lines.append("-" * 40)
        lines.append(f"Total Tests: {report.total_tests}")
        lines.append(f"Passed: {report.total_passed}")
        lines.append(f"Failed: {report.total_failed}")
        lines.append(f"Errors: {report.total_errors}")
        lines.append(f"Skipped: {report.total_skipped}")
        lines.append(f"Success Rate: {report.overall_success_rate:.1%}")
        lines.append("")
        
        # Workflow breakdown
        lines.append("WORKFLOW BREAKDOWN")
        lines.append("-" * 40)
        for suite in report.workflow_suites:
            lines.append(f"{suite.workflow_name.upper()}")
            lines.append(f"  Tests: {suite.total_tests}")
            lines.append(f"  Passed: {suite.passed_tests}")
            lines.append(f"  Failed: {suite.failed_tests}")
            lines.append(f"  Errors: {suite.error_tests}")
            lines.append(f"  Success Rate: {suite.success_rate:.1%}")
            lines.append(f"  Execution Time: {suite.total_execution_time:.2f}s")
            
            # Show failed tests
            failed_tests = [r for r in suite.test_results if r.status in [TestStatus.FAILED, TestStatus.ERROR]]
            if failed_tests:
                lines.append("  Failed Tests:")
                for test in failed_tests:
                    lines.append(f"    - {test.test_name}: {test.error_message}")
            lines.append("")
        
        # Environment info
        lines.append("ENVIRONMENT INFORMATION")
        lines.append("-" * 40)
        for key, value in report.environment_info.items():
            lines.append(f"{key}: {value}")
        lines.append("")
        
        return "\n".join(lines)
    
    def run_specific_workflow(self, workflow: str) -> WorkflowTestSuite:
        """Run tests for a specific workflow only."""
        if workflow not in self.workflow_tests:
            raise ValueError(f"Workflow '{workflow}' not found. Available: {list(self.workflow_tests.keys())}")
        
        log_dir = self.output_dir / "logs"
        log_dir.mkdir(exist_ok=True)
        
        with WorkflowTestLogger(f"{workflow}_{self.execution_id}", log_dir=str(log_dir)) as logger:
            self.logger = logger
            test_files = self.workflow_tests[workflow]
            return self._run_workflow_suite(workflow, test_files)
    
    def generate_test_visualization(self, report: TestExecutionReport) -> str:
        """Generate a simple ASCII visualization of test results."""
        lines = []
        lines.append("Test Results Visualization")
        lines.append("=" * 50)
        
        for suite in report.workflow_suites:
            # Create a simple bar chart
            total_width = 40
            passed_width = int((suite.passed_tests / suite.total_tests) * total_width) if suite.total_tests > 0 else 0
            failed_width = int((suite.failed_tests / suite.total_tests) * total_width) if suite.total_tests > 0 else 0
            error_width = int((suite.error_tests / suite.total_tests) * total_width) if suite.total_tests > 0 else 0
            remaining_width = total_width - passed_width - failed_width - error_width
            
            bar = "█" * passed_width + "▓" * failed_width + "░" * error_width + " " * remaining_width
            
            lines.append(f"{suite.workflow_name:20} [{bar}] {suite.success_rate:.1%}")
        
        lines.append("")
        lines.append("Legend: █ Passed  ▓ Failed  ░ Error")
        
        return "\n".join(lines)


def main():
    """Main entry point for the test runner."""
    parser = argparse.ArgumentParser(description="PrecipGen Workflow Test Runner")
    parser.add_argument("--workflows", nargs="+", help="Specific workflows to test")
    parser.add_argument("--output-dir", help="Output directory for test results")
    parser.add_argument("--parallel", action="store_true", help="Run workflows in parallel")
    parser.add_argument("--list-workflows", action="store_true", help="List available workflows")
    parser.add_argument("--config", help="Path to test configuration file")
    
    args = parser.parse_args()
    
    # Initialize runner
    config = TestConfig.from_file(args.config) if args.config else TestConfig()
    runner = WorkflowTestRunner(output_dir=args.output_dir, config=config)
    
    if args.list_workflows:
        print("Available workflows:")
        for workflow in runner.workflow_tests.keys():
            test_count = len(runner.workflow_tests[workflow])
            print(f"  {workflow}: {test_count} tests")
        return
    
    # Run tests
    try:
        report = runner.run_all_tests(workflows=args.workflows, parallel=args.parallel)
        
        # Print summary
        print("\n" + runner.generate_test_visualization(report))
        print(f"\nOverall Success Rate: {report.overall_success_rate:.1%}")
        
        # Exit with appropriate code
        sys.exit(0 if report.total_failed == 0 and report.total_errors == 0 else 1)
        
    except KeyboardInterrupt:
        print("\nTest execution interrupted by user")
        sys.exit(130)
    except Exception as e:
        print(f"Test execution failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
