"""
Specialized logging configuration for workflow testing framework.

This module extends PrecipGen's logging capabilities with testing-specific
features including test result logging, performance metrics, and organized
output management.
"""

import logging
import sys
from pathlib import Path
from typing import Optional, Dict, Any, List
from datetime import datetime
import json

from precipgen.utils.logging_config import PrecipGenLogger, get_logger


class WorkflowTestLogger:
    """
    Specialized logger for workflow testing framework.
    
    Extends PrecipGen's logging with testing-specific features:
    - Test result logging with structured format
    - Performance metrics collection
    - Test execution tracking
    - Organized log file management
    """
    
    def __init__(self, test_name: str, log_dir: Optional[str] = None,
                 log_level: str = "INFO"):
        """
        Initialize TestLogger for a specific test.
        
        Args:
            test_name: Name of the test or test suite
            log_dir: Optional directory for log files
            log_level: Logging level
        """
        self.test_name = test_name
        self.log_dir = Path(log_dir) if log_dir else Path("./test_logs")
        self.log_level = log_level
        
        # Create log directory
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        # Initialize base logger
        self.logger = get_logger(f"workflow_testing.{test_name}", log_level)
        
        # Set up test-specific file logging
        self.file_handler = None
        self._setup_test_logging()
        
        # Test execution tracking
        self.test_start_time: Optional[datetime] = None
        self.test_results: List[Dict[str, Any]] = []
        self.performance_metrics: Dict[str, Any] = {}
        
        self.logger.info(f"TestLogger initialized for test: {test_name}")
    
    def _setup_test_logging(self) -> None:
        """Set up test-specific file logging."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = self.log_dir / f"{self.test_name}_{timestamp}.log"
        
        # Create file handler for this test
        self.file_handler = logging.FileHandler(log_file)
        self.file_handler.setLevel(logging.DEBUG)  # File gets all messages
        
        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - "
            "%(funcName)s:%(lineno)d - %(message)s"
        )
        self.file_handler.setFormatter(formatter)
        
        # Add handler to logger
        self.logger.addHandler(self.file_handler)
        
        self.log_file_path = log_file
        self.logger.info(f"Test logging configured: {log_file}")
    
    def close(self) -> None:
        """Close the test logger and clean up file handlers."""
        if self.file_handler:
            self.logger.removeHandler(self.file_handler)
            self.file_handler.close()
            self.file_handler = None
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit with cleanup."""
        self.close()
    
    def start_test_execution(self, test_description: str,
                           parameters: Optional[Dict[str, Any]] = None) -> None:
        """
        Log the start of test execution.
        
        Args:
            test_description: Description of the test being executed
            parameters: Optional test parameters
        """
        self.test_start_time = datetime.now()
        
        log_msg = f"=== STARTING TEST EXECUTION: {test_description} ==="
        self.logger.info(log_msg)
        
        if parameters:
            param_str = ", ".join(f"{k}={v}" for k, v in parameters.items())
            self.logger.info(f"Test parameters: {param_str}")
        
        self.logger.info(f"Test start time: {self.test_start_time.isoformat()}")
    
    def end_test_execution(self, success: bool, summary: Optional[str] = None) -> None:
        """
        Log the end of test execution with results.
        
        Args:
            success: Whether the test execution was successful
            summary: Optional summary of results
        """
        end_time = datetime.now()
        
        if self.test_start_time:
            duration = (end_time - self.test_start_time).total_seconds()
            self.logger.info(f"Test execution duration: {duration:.3f} seconds")
        
        status = "SUCCESS" if success else "FAILURE"
        log_msg = f"=== TEST EXECUTION {status} ==="
        
        if success:
            self.logger.info(log_msg)
        else:
            self.logger.error(log_msg)
        
        if summary:
            self.logger.info(f"Test summary: {summary}")
        
        self.logger.info(f"Test end time: {end_time.isoformat()}")
    
    def log_test_result(self, test_case: str, success: bool, 
                       execution_time: float, details: Optional[Dict[str, Any]] = None) -> None:
        """
        Log individual test case result.
        
        Args:
            test_case: Name of the test case
            success: Whether the test case passed
            execution_time: Execution time in seconds
            details: Optional additional details
        """
        result = {
            "test_case": test_case,
            "success": success,
            "execution_time": execution_time,
            "timestamp": datetime.now().isoformat(),
            "details": details or {}
        }
        
        self.test_results.append(result)
        
        status = "PASS" if success else "FAIL"
        log_msg = f"Test case [{status}] {test_case} ({execution_time:.3f}s)"
        
        if success:
            self.logger.info(log_msg)
        else:
            self.logger.error(log_msg)
        
        if details:
            for key, value in details.items():
                self.logger.debug(f"  {key}: {value}")
    
    def log_performance_metric(self, metric_name: str, value: float,
                             unit: str = "", context: Optional[Dict[str, Any]] = None) -> None:
        """
        Log performance metric.
        
        Args:
            metric_name: Name of the performance metric
            value: Metric value
            unit: Optional unit of measurement
            context: Optional context information
        """
        if metric_name not in self.performance_metrics:
            self.performance_metrics[metric_name] = []
        
        metric_entry = {
            "value": value,
            "unit": unit,
            "timestamp": datetime.now().isoformat(),
            "context": context or {}
        }
        
        self.performance_metrics[metric_name].append(metric_entry)
        
        unit_str = f" {unit}" if unit else ""
        log_msg = f"Performance metric - {metric_name}: {value}{unit_str}"
        
        if context:
            context_str = ", ".join(f"{k}={v}" for k, v in context.items())
            log_msg += f" (context: {context_str})"
        
        self.logger.info(log_msg)
    
    def log_validation_result(self, validation_name: str, expected: Any,
                            actual: Any, passed: bool, threshold: Optional[float] = None) -> None:
        """
        Log validation result with expected vs actual comparison.
        
        Args:
            validation_name: Name of the validation
            expected: Expected value
            actual: Actual value
            passed: Whether validation passed
            threshold: Optional threshold value
        """
        status = "PASS" if passed else "FAIL"
        log_msg = f"Validation [{status}] {validation_name}"
        
        if passed:
            self.logger.info(log_msg)
        else:
            self.logger.error(log_msg)
        
        self.logger.info(f"  Expected: {expected}")
        self.logger.info(f"  Actual: {actual}")
        
        if threshold is not None:
            self.logger.info(f"  Threshold: {threshold}")
    
    def log_error_with_context(self, error: Exception, operation: str,
                             context: Optional[Dict[str, Any]] = None) -> None:
        """
        Log error with full context information.
        
        Args:
            error: Exception that occurred
            operation: Description of the operation that failed
            context: Additional context information
        """
        error_msg = f"Error in {operation}: {str(error)}"
        
        if hasattr(error, 'details') and error.details:
            error_msg += f" | Details: {error.details}"
        
        if context:
            error_msg += f" | Context: {context}"
        
        self.logger.error(error_msg, exc_info=True)
    
    def log_workflow_step(self, workflow: str, step: str, status: str,
                         details: Optional[Dict[str, Any]] = None) -> None:
        """
        Log workflow step execution.
        
        Args:
            workflow: Name of the workflow
            step: Name of the step
            status: Step status (started, completed, failed)
            details: Optional step details
        """
        log_msg = f"Workflow [{workflow}] Step [{step}] - {status.upper()}"
        
        if status.lower() == "failed":
            self.logger.error(log_msg)
        else:
            self.logger.info(log_msg)
        
        if details:
            for key, value in details.items():
                self.logger.debug(f"  {key}: {value}")
    
    def log_property_test_result(self, property_name: str, examples_run: int,
                               passed: bool, counterexample: Optional[Any] = None,
                               shrunk_example: Optional[Any] = None) -> None:
        """
        Log property-based test result.
        
        Args:
            property_name: Name of the property being tested
            examples_run: Number of examples that were run
            passed: Whether the property test passed
            counterexample: Optional counterexample that caused failure
            shrunk_example: Optional shrunk counterexample
        """
        status = "PASS" if passed else "FAIL"
        log_msg = f"Property test [{status}] {property_name} ({examples_run} examples)"
        
        if passed:
            self.logger.info(log_msg)
        else:
            self.logger.error(log_msg)
            
            if counterexample is not None:
                self.logger.error(f"  Counterexample: {counterexample}")
            
            if shrunk_example is not None:
                self.logger.error(f"  Shrunk example: {shrunk_example}")
    
    def generate_test_report(self) -> Dict[str, Any]:
        """
        Generate comprehensive test report.
        
        Returns:
            Dictionary containing test report data
        """
        total_tests = len(self.test_results)
        passed_tests = sum(1 for result in self.test_results if result["success"])
        failed_tests = total_tests - passed_tests
        
        success_rate = (passed_tests / total_tests * 100) if total_tests > 0 else 0.0
        
        total_execution_time = sum(result["execution_time"] for result in self.test_results)
        
        report = {
            "test_name": self.test_name,
            "timestamp": datetime.now().isoformat(),
            "summary": {
                "total_tests": total_tests,
                "passed_tests": passed_tests,
                "failed_tests": failed_tests,
                "success_rate": success_rate,
                "total_execution_time": total_execution_time
            },
            "test_results": self.test_results,
            "performance_metrics": self.performance_metrics,
            "log_file": str(self.log_file_path)
        }
        
        return report
    
    def save_test_report(self, output_file: Optional[str] = None) -> str:
        """
        Save test report to JSON file.
        
        Args:
            output_file: Optional output file path
            
        Returns:
            Path to saved report file
        """
        if output_file is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = self.log_dir / f"{self.test_name}_report_{timestamp}.json"
        
        report = self.generate_test_report()
        
        try:
            with open(output_file, 'w') as f:
                json.dump(report, f, indent=2)
            
            self.logger.info(f"Test report saved to: {output_file}")
            return str(output_file)
            
        except IOError as e:
            self.logger.error(f"Failed to save test report to {output_file}: {e}")
            raise
    
    def get_success_rate(self) -> float:
        """
        Get overall test success rate.
        
        Returns:
            Success rate as percentage (0-100)
        """
        if not self.test_results:
            return 0.0
        
        passed = sum(1 for result in self.test_results if result["success"])
        return (passed / len(self.test_results)) * 100
    
    def get_failed_tests(self) -> List[Dict[str, Any]]:
        """
        Get list of failed test results.
        
        Returns:
            List of failed test result dictionaries
        """
        return [result for result in self.test_results if not result["success"]]
    
    def get_performance_summary(self) -> Dict[str, Any]:
        """
        Get summary of performance metrics.
        
        Returns:
            Dictionary with performance metric summaries
        """
        summary = {}
        
        for metric_name, measurements in self.performance_metrics.items():
            values = [m["value"] for m in measurements]
            
            if values:
                summary[metric_name] = {
                    "count": len(values),
                    "min": min(values),
                    "max": max(values),
                    "mean": sum(values) / len(values),
                    "unit": measurements[0].get("unit", "")
                }
        
        return summary
