"""
Data Download Testing for PrecipGen Workflow Testing Framework.

This module tests GHCN data acquisition including:
- .dly file download and station data extraction to CSV
- Retry logic and network error handling  
- Organized file placement in station directories

Tests validate Requirements 2.6, 2.7 for the data acquisition workflow.
"""

import pytest
import tempfile
import shutil
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from unittest.mock import Mock, patch, MagicMock
import pandas as pd
import numpy as np
import requests

from tests.workflow_testing.infrastructure.directory_manager import DirectoryManager
from tests.workflow_testing.infrastructure.test_logger import WorkflowTestLogger
from precipgen.data.ghcn_downloader import GHCNDownloader
from precipgen.data.ghcn_parser import GHCNParser
from precipgen.utils.exceptions import NetworkError, DataLoadError


class DataDownloadTester:
    """
    Tester for GHCN data download and extraction functionality.
    
    Tests .dly file download, station data extraction to CSV, retry logic,
    network error handling, and organized file placement in station directories.
    """
    
    def __init__(self, cache_dir: Optional[str] = None):
        """
        Initialize data download tester.
        
        Args:
            cache_dir: Directory for caching test data
        """
        self.logger = WorkflowTestLogger("data_download_tester")
        self.cache_dir = cache_dir or tempfile.mkdtemp(prefix="precipgen_download_test_")
        self.directory_manager = DirectoryManager(base_dir=self.cache_dir)
        
        # Initialize components
        self.downloader = GHCNDownloader(cache_dir=self.cache_dir)
        # Note: GHCNParser will be initialized with specific file paths when needed
        
    def test_dly_file_download(self, station_id: str) -> bool:
        """
        Test downloading .dly file for a specific station.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if download successful, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Download .dly file
            file_path = self.downloader.download_station_data(station_id)
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            # Verify file exists and has content
            file_exists = Path(file_path).exists()
            file_size = Path(file_path).stat().st_size if file_exists else 0
            
            success = file_exists and file_size > 0
            
            self.logger.log_test_result(
                f"dly_download_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "file_path": file_path,
                    "file_exists": file_exists,
                    "file_size_bytes": file_size
                }
            )
            
            return success
            
        except (NetworkError, DataLoadError) as e:
            self.logger.log_test_result(
                f"dly_download_{station_id}_error",
                False,
                0.0,
                details={"station_id": station_id, "error": str(e)}
            )
            return False
        except Exception as e:
            self.logger.log_error_with_context(e, f"dly_download_{station_id}")
            return False
    
    def test_station_data_extraction_to_csv(self, station_id: str) -> bool:
        """
        Test extracting station data from .dly file to CSV format.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if extraction successful, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # First download the .dly file
            dly_file_path = self.downloader.download_station_data(station_id)
            
            # Create station directory for organized output
            station_dir = self.directory_manager.create_station_directory(station_id)
            data_dir = station_dir / "data"
            
            # Extract precipitation data to CSV
            parser = GHCNParser(dly_file_path)
            precipitation_data = parser.extract_precipitation_data()
            
            # Save to organized directory structure
            csv_file_path = data_dir / "precipitation.csv"
            precipitation_data.to_csv(csv_file_path, index=True)
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            # Verify CSV file was created and has data
            csv_exists = csv_file_path.exists()
            csv_size = csv_file_path.stat().st_size if csv_exists else 0
            data_rows = len(precipitation_data) if precipitation_data is not None else 0
            
            success = csv_exists and csv_size > 0 and data_rows > 0
            
            self.logger.log_test_result(
                f"csv_extraction_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "dly_file": dly_file_path,
                    "csv_file": str(csv_file_path),
                    "csv_exists": csv_exists,
                    "csv_size_bytes": csv_size,
                    "data_rows": data_rows
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"csv_extraction_{station_id}")
            return False
    
    def test_organized_file_placement(self, station_id: str) -> bool:
        """
        Test that files are placed in organized directory structure.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if files are properly organized, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Download and extract data
            dly_file_path = self.downloader.download_station_data(station_id)
            station_dir = self.directory_manager.create_station_directory(station_id)
            data_dir = station_dir / "data"
            
            # Extract and save data
            parser = GHCNParser(dly_file_path)
            precipitation_data = parser.extract_precipitation_data()
            csv_file_path = data_dir / "precipitation.csv"
            precipitation_data.to_csv(csv_file_path, index=True)
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            # Verify organized directory structure
            expected_structure = {
                "station_dir_exists": station_dir.exists(),
                "data_dir_exists": data_dir.exists(),
                "csv_in_data_dir": csv_file_path.exists(),
                "correct_station_name": station_dir.name == station_id
            }
            
            success = all(expected_structure.values())
            
            self.logger.log_test_result(
                f"organized_placement_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "station_dir": str(station_dir),
                    "structure_check": expected_structure
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"organized_placement_{station_id}")
            return False
    
    def test_retry_logic_with_network_errors(self, station_id: str) -> bool:
        """
        Test retry logic when network errors occur.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if retry logic works correctly, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Mock network errors for first few attempts
            original_get = requests.get
            call_count = 0
            
            def mock_get_with_failures(*args, **kwargs):
                nonlocal call_count
                call_count += 1
                
                # Fail first 2 attempts, succeed on 3rd
                if call_count <= 2:
                    raise requests.ConnectionError("Simulated network error")
                else:
                    return original_get(*args, **kwargs)
            
            with patch('requests.get', side_effect=mock_get_with_failures):
                try:
                    # This should succeed after retries
                    file_path = self.downloader.download_station_data(station_id)
                    success = Path(file_path).exists()
                    retry_worked = call_count > 1  # Should have retried
                    
                except NetworkError:
                    # If it still fails, that's also valid behavior
                    success = True  # The retry logic attempted retries
                    retry_worked = call_count > 1
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            self.logger.log_test_result(
                f"retry_logic_{station_id}",
                success and retry_worked,
                execution_time,
                details={
                    "station_id": station_id,
                    "retry_attempts": call_count,
                    "retry_worked": retry_worked,
                    "final_success": success
                }
            )
            
            return success and retry_worked
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"retry_logic_{station_id}")
            return False
    
    def test_network_error_handling(self, station_id: str) -> bool:
        """
        Test proper handling of network errors.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if errors are handled gracefully, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Mock persistent network failure
            def mock_get_failure(*args, **kwargs):
                raise requests.ConnectionError("Persistent network error")
            
            with patch('requests.get', side_effect=mock_get_failure):
                try:
                    file_path = self.downloader.download_station_data(station_id)
                    # Should not reach here
                    success = False
                    error_handled = False
                    
                except NetworkError as e:
                    # Expected error type
                    success = True
                    error_handled = True
                    error_message = str(e)
                    
                    # Verify error message is informative
                    informative = any(keyword in error_message.lower() for keyword in 
                                    ['network', 'connection', 'download', 'failed'])
                    success = success and informative
                    
                except Exception as e:
                    # Unexpected error type
                    success = False
                    error_handled = False
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            self.logger.log_test_result(
                f"network_error_handling_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "error_handled": error_handled,
                    "proper_error_type": success
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"network_error_handling_{station_id}")
            return False
    
    def test_multiple_station_downloads(self, station_ids: List[str]) -> Dict[str, bool]:
        """
        Test downloading data for multiple stations.
        
        Args:
            station_ids: List of GHCN station identifiers
            
        Returns:
            Dictionary mapping station_id -> success status
        """
        results = {}
        
        self.logger.start_test_execution(f"Multiple station downloads: {len(station_ids)} stations")
        
        try:
            for i, station_id in enumerate(station_ids):
                try:
                    # Test download and extraction for each station
                    download_success = self.test_dly_file_download(station_id)
                    extraction_success = self.test_station_data_extraction_to_csv(station_id)
                    organization_success = self.test_organized_file_placement(station_id)
                    
                    # Overall success for this station
                    station_success = download_success and extraction_success and organization_success
                    results[station_id] = station_success
                    
                    self.logger.log_test_result(
                        f"multi_download_{station_id}",
                        station_success,
                        0.0,
                        details={
                            "station_id": station_id,
                            "progress": f"{i+1}/{len(station_ids)}",
                            "download": download_success,
                            "extraction": extraction_success,
                            "organization": organization_success
                        }
                    )
                    
                except Exception as e:
                    results[station_id] = False
                    self.logger.log_error_with_context(e, f"multi_download_{station_id}")
            
            overall_success = all(results.values())
            success_rate = sum(results.values()) / len(results) if results else 0.0
            
            self.logger.end_test_execution(
                overall_success,
                f"Multiple downloads completed: {success_rate:.1%} success rate"
            )
            
            return results
            
        except Exception as e:
            self.logger.log_error_with_context(e, "multiple_station_downloads")
            self.logger.end_test_execution(False, f"Multiple downloads failed: {e}")
            return results
    
    def run_comprehensive_tests(self) -> Dict[str, bool]:
        """
        Run comprehensive data download tests.
        
        Returns:
            Dictionary of test results
        """
        results = {}
        
        self.logger.start_test_execution("Comprehensive Data Download Tests")
        
        try:
            # Test with Salt Lake City station (gold standard)
            slc_station = "USC00426128"
            
            # Individual component tests
            results["dly_download"] = self.test_dly_file_download(slc_station)
            results["csv_extraction"] = self.test_station_data_extraction_to_csv(slc_station)
            results["organized_placement"] = self.test_organized_file_placement(slc_station)
            results["retry_logic"] = self.test_retry_logic_with_network_errors(slc_station)
            results["error_handling"] = self.test_network_error_handling(slc_station)
            
            # Multiple station test
            test_stations = ["USC00426128", "USW00024127", "USC00305426"]
            multi_results = self.test_multiple_station_downloads(test_stations)
            results["multiple_downloads"] = all(multi_results.values())
            
            overall_success = all(results.values())
            
            self.logger.end_test_execution(
                overall_success,
                f"Completed {len(results)} tests, {sum(results.values())} passed"
            )
            
            return results
            
        except Exception as e:
            self.logger.log_error_with_context(e, "comprehensive_tests")
            self.logger.end_test_execution(False, f"Comprehensive tests failed: {e}")
            return results


class TestDataDownloadWorkflow:
    """Test suite for data download workflow components."""
    
    def setup_method(self):
        """Set up test environment."""
        self.temp_dir = tempfile.mkdtemp(prefix="precipgen_data_download_test_")
        self.tester = DataDownloadTester(cache_dir=self.temp_dir)
        
    def teardown_method(self):
        """Clean up test environment."""
        if hasattr(self, 'temp_dir') and Path(self.temp_dir).exists():
            shutil.rmtree(self.temp_dir)
    
    def test_data_download_tester_initialization(self):
        """Test that DataDownloadTester initializes correctly."""
        assert self.tester is not None
        assert self.tester.cache_dir == self.temp_dir
        assert self.tester.downloader is not None
        # Note: parser is initialized per file, not as instance variable
        assert self.tester.directory_manager is not None
    
    def test_dly_file_download_salt_lake_city(self):
        """Test downloading .dly file for Salt Lake City (gold standard)."""
        # Use Salt Lake City station
        slc_station = "USC00426128"
        
        success = self.tester.test_dly_file_download(slc_station)
        assert success == True, "Should successfully download Salt Lake City .dly file"
        
        # Verify file exists in cache
        expected_file = Path(self.temp_dir) / f"{slc_station}.dly"
        assert expected_file.exists(), "Downloaded .dly file should exist in cache"
        assert expected_file.stat().st_size > 0, ".dly file should have content"
    
    def test_csv_extraction_salt_lake_city(self):
        """Test extracting CSV data for Salt Lake City."""
        slc_station = "USC00426128"
        
        success = self.tester.test_station_data_extraction_to_csv(slc_station)
        assert success == True, "Should successfully extract CSV data"
        
        # Verify organized directory structure
        station_dir = Path(self.temp_dir) / slc_station
        csv_file = station_dir / "data" / "precipitation.csv"
        
        assert station_dir.exists(), "Station directory should exist"
        assert csv_file.exists(), "CSV file should exist in data subdirectory"
        assert csv_file.stat().st_size > 0, "CSV file should have content"
    
    def test_organized_file_placement(self):
        """Test that files are placed in organized directory structure."""
        slc_station = "USC00426128"
        
        success = self.tester.test_organized_file_placement(slc_station)
        assert success == True, "Files should be organized correctly"
        
        # Verify complete directory structure
        station_dir = Path(self.temp_dir) / slc_station
        data_dir = station_dir / "data"
        
        assert station_dir.exists(), "Station directory should exist"
        assert data_dir.exists(), "Data subdirectory should exist"
        assert station_dir.name == slc_station, "Directory name should match station ID"
    
    def test_network_error_handling(self):
        """Test network error handling."""
        slc_station = "USC00426128"
        
        success = self.tester.test_network_error_handling(slc_station)
        assert success == True, "Should handle network errors gracefully"
    
    def test_comprehensive_workflow(self):
        """Test the complete data download workflow."""
        results = self.tester.run_comprehensive_tests()
        
        # Should have multiple test results
        assert len(results) > 0, "Should have test results"
        
        # Most tests should pass (allow for some network-dependent failures)
        success_rate = sum(results.values()) / len(results)
        assert success_rate >= 0.6, f"Success rate too low: {success_rate:.2f}"
