"""
Batch Processing Testing for PrecipGen Workflows.

This module tests batch processing with multiple stations, implements graceful
error handling for individual failures, and tests organized output generation
for each station.

Tests validate Requirement 8.4 for batch processing testing.
"""

import pytest
import tempfile
import shutil
import time
import threading
from pathlib import Path
from typing import List, Dict, Optional, Tuple, Union, Any
from unittest.mock import Mock, patch, MagicMock
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
from hypothesis import given, strategies as st, assume, settings, HealthCheck

from tests.workflow_testing.infrastructure.test_logger import WorkflowTestLogger
from tests.workflow_testing.infrastructure.test_config import TestConfig
from tests.workflow_testing.test_integration_workflow import EndToEndTester, WorkflowIndependenceTester
from precipgen.engines.analytical import AnalyticalEngine, ParameterManifest
from precipgen.engines.simulation import SimulationEngine
from precipgen.convenience.validation_suite import ValidationSuite
from precipgen.utils.exceptions import PrecipGenError, ValidationError, SimulationEngineError
from precipgen.utils.directory_manager import DirectoryManager


@dataclass
class BatchProcessingResult:
    """Result from batch processing multiple locations."""
    total_locations: int
    successful_locations: int
    failed_locations: int
    success_rate: float
    execution_time: float
    location_results: Dict[str, Dict[str, Any]]
    error_summary: Dict[str, List[str]]
    output_directories: List[str]


@dataclass
class LocationProcessingResult:
    """Result from processing a single location in batch."""
    location_id: str
    coordinates: Tuple[float, float]
    location_name: Optional[str]
    success: bool
    execution_time: float
    station_id: Optional[str] = None
    validation_score: Optional[float] = None
    error_message: Optional[str] = None
    output_directory: Optional[str] = None
    workflow_results: Dict[str, Any] = field(default_factory=dict)


class BatchTester:
    """
    Test batch processing with multiple stations and graceful error handling.
    
    Validates Requirement 8.4 for batch processing testing.
    """
    
    def __init__(self, base_output_dir: Optional[Path] = None, max_workers: int = 4):
        """
        Initialize batch tester.
        
        Args:
            base_output_dir: Base directory for all outputs
            max_workers: Maximum number of parallel workers
        """
        self.base_output_dir = base_output_dir or Path("./precipgen_output")
        self.directory_manager = DirectoryManager(self.base_output_dir)
        self.logger = WorkflowTestLogger("BatchTester")
        self.config = TestConfig()
        self.max_workers = max_workers
        
        # Error handling configuration
        self.continue_on_error = True
        self.retry_failed_locations = False
        self.max_retries = 1
    
    def test_batch_processing_multiple_locations(self, locations: List[Tuple[Tuple[float, float], str]]) -> BatchProcessingResult:
        """
        Test batch processing with multiple locations.
        
        Args:
            locations: List of ((lat, lon), location_name) tuples
            
        Returns:
            BatchProcessingResult with comprehensive results
        """
        start_time = time.time()
        
        self.logger.logger.info(f"Starting batch processing for {len(locations)} locations")
        
        location_results = {}
        error_summary = {
            'station_discovery_errors': [],
            'data_acquisition_errors': [],
            'analysis_errors': [],
            'simulation_errors': [],
            'validation_errors': [],
            'general_errors': []
        }
        output_directories = []
        
        successful_count = 0
        failed_count = 0
        
        # Process locations sequentially for better error tracking
        # In production, this could use parallel processing
        for i, (coordinates, location_name) in enumerate(locations):
            location_id = f"location_{i+1}_{location_name.replace(' ', '_').replace(',', '')}"
            
            try:
                self.logger.logger.info(f"Processing location {i+1}/{len(locations)}: {location_name} {coordinates}")
                
                result = self._process_single_location(coordinates, location_name, location_id)
                location_results[location_id] = result.__dict__
                
                if result.success:
                    successful_count += 1
                    if result.output_directory:
                        output_directories.append(result.output_directory)
                else:
                    failed_count += 1
                    self._categorize_error(result.error_message, error_summary)
                
            except Exception as e:
                failed_count += 1
                error_result = LocationProcessingResult(
                    location_id=location_id,
                    coordinates=coordinates,
                    location_name=location_name,
                    success=False,
                    execution_time=0.0,
                    error_message=str(e)
                )
                location_results[location_id] = error_result.__dict__
                error_summary['general_errors'].append(f"{location_name}: {str(e)}")
                
                self.logger.logger.error(f"Failed to process location {location_name}: {e}")
                
                # Continue processing other locations if configured to do so
                if not self.continue_on_error:
                    break
        
        execution_time = time.time() - start_time
        success_rate = successful_count / len(locations) if locations else 0.0
        
        batch_result = BatchProcessingResult(
            total_locations=len(locations),
            successful_locations=successful_count,
            failed_locations=failed_count,
            success_rate=success_rate,
            execution_time=execution_time,
            location_results=location_results,
            error_summary=error_summary,
            output_directories=output_directories
        )
        
        self.logger.logger.info(
            f"Batch processing completed: {successful_count}/{len(locations)} successful "
            f"({success_rate:.1%} success rate) in {execution_time:.1f}s"
        )
        
        return batch_result
    
    def test_parallel_batch_processing(self, locations: List[Tuple[Tuple[float, float], str]]) -> BatchProcessingResult:
        """
        Test parallel batch processing with multiple workers.
        
        Args:
            locations: List of ((lat, lon), location_name) tuples
            
        Returns:
            BatchProcessingResult with comprehensive results
        """
        start_time = time.time()
        
        self.logger.logger.info(f"Starting parallel batch processing for {len(locations)} locations with {self.max_workers} workers")
        
        location_results = {}
        error_summary = {
            'station_discovery_errors': [],
            'data_acquisition_errors': [],
            'analysis_errors': [],
            'simulation_errors': [],
            'validation_errors': [],
            'general_errors': []
        }
        output_directories = []
        
        successful_count = 0
        failed_count = 0
        
        # Use ThreadPoolExecutor for parallel processing
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all tasks
            future_to_location = {}
            for i, (coordinates, location_name) in enumerate(locations):
                location_id = f"location_{i+1}_{location_name.replace(' ', '_').replace(',', '')}"
                future = executor.submit(self._process_single_location, coordinates, location_name, location_id)
                future_to_location[future] = (location_id, coordinates, location_name)
            
            # Collect results as they complete
            for future in as_completed(future_to_location):
                location_id, coordinates, location_name = future_to_location[future]
                
                try:
                    result = future.result()
                    location_results[location_id] = result.__dict__
                    
                    if result.success:
                        successful_count += 1
                        if result.output_directory:
                            output_directories.append(result.output_directory)
                    else:
                        failed_count += 1
                        self._categorize_error(result.error_message, error_summary)
                    
                    self.logger.logger.info(f"Completed processing {location_name}: {'SUCCESS' if result.success else 'FAILED'}")
                    
                except Exception as e:
                    failed_count += 1
                    error_result = LocationProcessingResult(
                        location_id=location_id,
                        coordinates=coordinates,
                        location_name=location_name,
                        success=False,
                        execution_time=0.0,
                        error_message=str(e)
                    )
                    location_results[location_id] = error_result.__dict__
                    error_summary['general_errors'].append(f"{location_name}: {str(e)}")
                    
                    self.logger.logger.error(f"Failed to process location {location_name}: {e}")
        
        execution_time = time.time() - start_time
        success_rate = successful_count / len(locations) if locations else 0.0
        
        batch_result = BatchProcessingResult(
            total_locations=len(locations),
            successful_locations=successful_count,
            failed_locations=failed_count,
            success_rate=success_rate,
            execution_time=execution_time,
            location_results=location_results,
            error_summary=error_summary,
            output_directories=output_directories
        )
        
        self.logger.logger.info(
            f"Parallel batch processing completed: {successful_count}/{len(locations)} successful "
            f"({success_rate:.1%} success rate) in {execution_time:.1f}s"
        )
        
        return batch_result
    
    def test_graceful_error_handling(self, locations: List[Tuple[Tuple[float, float], str]]) -> Dict[str, Any]:
        """
        Test graceful error handling for individual failures.
        
        Args:
            locations: List of locations including some that should fail
            
        Returns:
            Dictionary with error handling test results
        """
        self.logger.logger.info("Testing graceful error handling with mixed valid/invalid locations")
        
        # Add some intentionally problematic locations
        test_locations = locations.copy()
        
        # Add invalid coordinates
        test_locations.append(((200.0, -500.0), "Invalid Coordinates"))
        
        # Add coordinates in ocean (no nearby stations)
        test_locations.append(((0.0, 0.0), "Middle of Ocean"))
        
        # Process with error handling
        original_continue_setting = self.continue_on_error
        self.continue_on_error = True
        
        try:
            batch_result = self.test_batch_processing_multiple_locations(test_locations)
            
            # Analyze error handling effectiveness
            error_handling_results = {
                'total_locations_tested': len(test_locations),
                'locations_with_errors': batch_result.failed_locations,
                'continued_after_errors': batch_result.successful_locations > 0,
                'error_categories': batch_result.error_summary,
                'graceful_degradation': True,  # If we get here, degradation was graceful
                'output_directories_preserved': len(batch_result.output_directories),
                'success_rate': batch_result.success_rate
            }
            
            # Verify that successful locations still produced outputs
            successful_outputs_valid = True
            for output_dir in batch_result.output_directories:
                if not Path(output_dir).exists():
                    successful_outputs_valid = False
                    break
            
            error_handling_results['successful_outputs_preserved'] = successful_outputs_valid
            
            self.logger.logger.info(
                f"Error handling test completed: {batch_result.successful_locations} successful, "
                f"{batch_result.failed_locations} failed, graceful degradation: {error_handling_results['graceful_degradation']}"
            )
            
            return error_handling_results
            
        finally:
            self.continue_on_error = original_continue_setting
    
    def test_organized_output_generation(self, locations: List[Tuple[Tuple[float, float], str]]) -> Dict[str, Any]:
        """
        Test organized output generation for each station.
        
        Args:
            locations: List of locations to test
            
        Returns:
            Dictionary with output organization test results
        """
        self.logger.logger.info("Testing organized output generation for multiple stations")
        
        batch_result = self.test_batch_processing_multiple_locations(locations)
        
        # Analyze output organization
        organization_results = {
            'total_output_directories': len(batch_result.output_directories),
            'directory_structure_valid': True,
            'required_files_present': True,
            'station_isolation_maintained': True,
            'directory_details': []
        }
        
        for output_dir in batch_result.output_directories:
            output_path = Path(output_dir)
            station_id = output_path.name
            
            # Check directory structure
            directory_analysis = self._analyze_output_directory(output_path)
            organization_results['directory_details'].append(directory_analysis)
            
            # Update overall results
            if not directory_analysis['structure_valid']:
                organization_results['directory_structure_valid'] = False
            
            if not directory_analysis['required_files_present']:
                organization_results['required_files_present'] = False
        
        # Check for cross-contamination between stations
        station_ids = [Path(d).name for d in batch_result.output_directories]
        unique_station_ids = set(station_ids)
        
        if len(unique_station_ids) != len(station_ids):
            organization_results['station_isolation_maintained'] = False
        
        self.logger.logger.info(
            f"Output organization test completed: {len(batch_result.output_directories)} directories, "
            f"structure valid: {organization_results['directory_structure_valid']}, "
            f"isolation maintained: {organization_results['station_isolation_maintained']}"
        )
        
        return organization_results
    
    def _process_single_location(self, coordinates: Tuple[float, float], 
                               location_name: str, location_id: str) -> LocationProcessingResult:
        """
        Process a single location through the complete workflow.
        
        Args:
            coordinates: Location coordinates
            location_name: Human-readable location name
            location_id: Unique identifier for this location
            
        Returns:
            LocationProcessingResult with processing details
        """
        start_time = time.time()
        
        try:
            # Create end-to-end tester for this location
            # Use a unique cache directory to avoid conflicts
            cache_dir = tempfile.mkdtemp(prefix=f"batch_test_{location_id}_")
            
            end_to_end_tester = EndToEndTester(cache_dir=cache_dir)
            
            # Execute complete workflow
            workflow_result = end_to_end_tester.test_end_to_end_workflow_with_coordinates(
                coordinates, location_name
            )
            
            execution_time = time.time() - start_time
            
            if workflow_result["success"]:
                return LocationProcessingResult(
                    location_id=location_id,
                    coordinates=coordinates,
                    location_name=location_name,
                    success=True,
                    execution_time=execution_time,
                    station_id=workflow_result.get("station_id"),
                    validation_score=workflow_result.get("validation_score"),
                    output_directory=workflow_result.get("output_directory"),
                    workflow_results=workflow_result.get("workflow_results", {})
                )
            else:
                return LocationProcessingResult(
                    location_id=location_id,
                    coordinates=coordinates,
                    location_name=location_name,
                    success=False,
                    execution_time=execution_time,
                    error_message=workflow_result.get("error", "Unknown workflow error")
                )
        
        except Exception as e:
            execution_time = time.time() - start_time
            return LocationProcessingResult(
                location_id=location_id,
                coordinates=coordinates,
                location_name=location_name,
                success=False,
                execution_time=execution_time,
                error_message=str(e)
            )
    
    def _categorize_error(self, error_message: Optional[str], error_summary: Dict[str, List[str]]) -> None:
        """Categorize error message into appropriate error type."""
        if not error_message:
            return
        
        error_lower = error_message.lower()
        
        if any(keyword in error_lower for keyword in ['station', 'discovery', 'geocoding']):
            error_summary['station_discovery_errors'].append(error_message)
        elif any(keyword in error_lower for keyword in ['data', 'download', 'ghcn']):
            error_summary['data_acquisition_errors'].append(error_message)
        elif any(keyword in error_lower for keyword in ['analysis', 'parameter', 'analytical']):
            error_summary['analysis_errors'].append(error_message)
        elif any(keyword in error_lower for keyword in ['simulation', 'synthetic', 'wgen']):
            error_summary['simulation_errors'].append(error_message)
        elif any(keyword in error_lower for keyword in ['validation', 'validate']):
            error_summary['validation_errors'].append(error_message)
        else:
            error_summary['general_errors'].append(error_message)
    
    def _analyze_output_directory(self, output_path: Path) -> Dict[str, Any]:
        """
        Analyze output directory structure and contents.
        
        Args:
            output_path: Path to station output directory
            
        Returns:
            Dictionary with analysis results
        """
        analysis = {
            'station_id': output_path.name,
            'directory_path': str(output_path),
            'structure_valid': True,
            'required_files_present': True,
            'subdirectories': {},
            'total_files': 0,
            'missing_components': []
        }
        
        # Check for required subdirectories
        required_subdirs = ['data', 'analysis', 'simulation', 'validation']
        
        for subdir in required_subdirs:
            subdir_path = output_path / subdir
            
            if subdir_path.exists():
                files = list(subdir_path.rglob('*'))
                file_count = len([f for f in files if f.is_file()])
                
                analysis['subdirectories'][subdir] = {
                    'exists': True,
                    'file_count': file_count,
                    'files': [str(f.relative_to(subdir_path)) for f in files if f.is_file()]
                }
                analysis['total_files'] += file_count
            else:
                analysis['subdirectories'][subdir] = {
                    'exists': False,
                    'file_count': 0,
                    'files': []
                }
                analysis['structure_valid'] = False
                analysis['missing_components'].append(f"Missing {subdir} directory")
        
        # Check for key files that indicate workflow completion
        required_files = {
            'data': ['precipitation.csv', 'quality_report.json'],
            'analysis': ['params.json'],
            'simulation': ['synthetic.csv'],
            'validation': ['validation_report.json']
        }
        
        for workflow, files in required_files.items():
            if workflow in analysis['subdirectories'] and analysis['subdirectories'][workflow]['exists']:
                workflow_files = analysis['subdirectories'][workflow]['files']
                
                for required_file in files:
                    if not any(required_file in f for f in workflow_files):
                        analysis['required_files_present'] = False
                        analysis['missing_components'].append(f"Missing {required_file} in {workflow}")
        
        return analysis


# Unit tests for BatchTester functionality
class TestBatchTesterUnit:
    """Unit tests for BatchTester functionality."""
    
    def setup_method(self):
        """Set up test environment."""
        self.temp_dir = Path(tempfile.mkdtemp())
        self.batch_tester = BatchTester(self.temp_dir, max_workers=2)
    
    def teardown_method(self):
        """Clean up test environment."""
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
    
    def test_batch_tester_initialization(self):
        """Test BatchTester initialization."""
        assert self.batch_tester is not None
        assert self.batch_tester.base_output_dir == self.temp_dir
        assert self.batch_tester.max_workers == 2
        assert self.batch_tester.continue_on_error is True
    
    def test_location_processing_result_creation(self):
        """Test LocationProcessingResult creation."""
        result = LocationProcessingResult(
            location_id="test_location",
            coordinates=(40.0, -111.0),
            location_name="Test Location",
            success=True,
            execution_time=1.5,
            station_id="TEST123",
            validation_score=0.85
        )
        
        assert result.location_id == "test_location"
        assert result.coordinates == (40.0, -111.0)
        assert result.location_name == "Test Location"
        assert result.success is True
        assert result.execution_time == 1.5
        assert result.station_id == "TEST123"
        assert result.validation_score == 0.85
    
    def test_batch_processing_result_creation(self):
        """Test BatchProcessingResult creation."""
        result = BatchProcessingResult(
            total_locations=5,
            successful_locations=4,
            failed_locations=1,
            success_rate=0.8,
            execution_time=30.0,
            location_results={},
            error_summary={},
            output_directories=[]
        )
        
        assert result.total_locations == 5
        assert result.successful_locations == 4
        assert result.failed_locations == 1
        assert result.success_rate == 0.8
        assert result.execution_time == 30.0
    
    def test_error_categorization(self):
        """Test error message categorization."""
        error_summary = {
            'station_discovery_errors': [],
            'data_acquisition_errors': [],
            'analysis_errors': [],
            'simulation_errors': [],
            'validation_errors': [],
            'general_errors': []
        }
        
        # Test different error types
        self.batch_tester._categorize_error("Station discovery failed", error_summary)
        self.batch_tester._categorize_error("Data download error", error_summary)
        self.batch_tester._categorize_error("Analysis parameter extraction failed", error_summary)
        self.batch_tester._categorize_error("Simulation engine error", error_summary)
        self.batch_tester._categorize_error("Validation failed", error_summary)
        self.batch_tester._categorize_error("Unknown error", error_summary)
        
        assert len(error_summary['station_discovery_errors']) == 1
        assert len(error_summary['data_acquisition_errors']) == 1
        assert len(error_summary['analysis_errors']) == 1
        assert len(error_summary['simulation_errors']) == 1
        assert len(error_summary['validation_errors']) == 1
        assert len(error_summary['general_errors']) == 1
    
    def test_output_directory_analysis(self):
        """Test output directory structure analysis."""
        # Create a test station directory structure
        station_id = "TEST123456"
        station_dir = self.temp_dir / station_id
        
        # Create subdirectories
        for subdir in ['data', 'analysis', 'simulation', 'validation']:
            (station_dir / subdir).mkdir(parents=True)
        
        # Create some test files
        (station_dir / 'data' / 'precipitation.csv').touch()
        (station_dir / 'data' / 'quality_report.json').touch()
        (station_dir / 'analysis' / 'params.json').touch()
        (station_dir / 'simulation' / 'synthetic.csv').touch()
        (station_dir / 'validation' / 'validation_report.json').touch()
        
        # Analyze the directory
        analysis = self.batch_tester._analyze_output_directory(station_dir)
        
        assert analysis['station_id'] == station_id
        assert analysis['structure_valid'] is True
        assert analysis['required_files_present'] is True
        assert analysis['total_files'] == 5
        assert len(analysis['missing_components']) == 0
        
        # Test with missing files
        (station_dir / 'data' / 'precipitation.csv').unlink()
        analysis_missing = self.batch_tester._analyze_output_directory(station_dir)
        
        assert analysis_missing['required_files_present'] is False
        assert len(analysis_missing['missing_components']) > 0


# Property-based tests
class TestBatchProcessingProperties:
    """Property-based tests for batch processing functionality."""
    
    def setup_method(self):
        """Set up test environment."""
        self.temp_dir = Path(tempfile.mkdtemp())
        self.batch_tester = BatchTester(self.temp_dir, max_workers=2)
    
    def teardown_method(self):
        """Clean up test environment."""
        if self.temp_dir.exists():
            shutil.rmtree(self.temp_dir)
    
    @given(
        locations=st.lists(
            st.tuples(
                st.tuples(
                    st.floats(min_value=-90, max_value=90),  # latitude
                    st.floats(min_value=-180, max_value=180)  # longitude
                ),
                st.text(min_size=1, max_size=20, alphabet=st.characters(whitelist_categories=('Lu', 'Ll', 'Nd', 'Pc')))
            ),
            min_size=1,
            max_size=3  # Keep small for testing
        )
    )
    @settings(max_examples=10, deadline=60000, suppress_health_check=[HealthCheck.too_slow])
    def test_property_batch_processing_graceful_degradation(self, locations):
        """
        Property test: Batch processing should handle errors gracefully and continue processing.
        
        **Property**: For any list of locations, batch processing should complete without
        crashing and provide meaningful results even if individual locations fail.
        **Validates: Requirements 8.4**
        """
        assume(len(locations) > 0)
        
        # Enable graceful error handling
        self.batch_tester.continue_on_error = True
        
        try:
            # Process the locations
            result = self.batch_tester.test_batch_processing_multiple_locations(locations)
            
            # Verify graceful degradation properties
            assert isinstance(result, BatchProcessingResult)
            assert result.total_locations == len(locations)
            assert result.successful_locations + result.failed_locations == result.total_locations
            assert 0.0 <= result.success_rate <= 1.0
            assert result.execution_time >= 0.0
            assert isinstance(result.location_results, dict)
            assert isinstance(result.error_summary, dict)
            assert isinstance(result.output_directories, list)
            
            # Each location should have a result
            assert len(result.location_results) == len(locations)
            
            # Success rate should match counts
            expected_success_rate = result.successful_locations / result.total_locations
            assert abs(result.success_rate - expected_success_rate) < 0.001
            
            # Output directories should match successful locations
            assert len(result.output_directories) == result.successful_locations
            
        except Exception as e:
            # Even if processing fails, it should fail gracefully
            pytest.fail(f"Batch processing should not crash, but got: {e}")
