"""
Quality Assessment Testing for PrecipGen Workflow Testing Framework.

This module tests data quality analysis including:
- Quality report generation and validation
- Station metadata extraction and storage
- Quality threshold validation

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

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

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.convenience.data_quality import DataQualityAnalyzer, DataQualityMetrics
from precipgen.convenience.models import StationInfo


class QualityAssessmentTester:
    """
    Tester for data quality assessment functionality.
    
    Tests quality report generation, station metadata extraction and storage,
    and quality threshold validation for precipitation data.
    """
    
    def __init__(self, cache_dir: Optional[str] = None):
        """
        Initialize quality assessment tester.
        
        Args:
            cache_dir: Directory for caching test data
        """
        self.logger = WorkflowTestLogger("quality_assessment_tester")
        self.cache_dir = cache_dir or tempfile.mkdtemp(prefix="precipgen_quality_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
        self.quality_analyzer = DataQualityAnalyzer()
        
    def test_quality_report_generation(self, station_id: str) -> bool:
        """
        Test generating quality report for station data.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if quality report generated successfully, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Download and parse data
            dly_file_path = self.downloader.download_station_data(station_id)
            parser = GHCNParser(dly_file_path)
            precipitation_data = parser.extract_precipitation_data()
            
            # Generate quality report
            quality_metrics = self.quality_analyzer.analyze_station_quality(station_id, dly_file_path)
            
            # Create station directory for organized output
            station_dir = self.directory_manager.create_station_directory(station_id)
            data_dir = station_dir / "data"
            
            # Save quality report
            quality_report_path = data_dir / "quality_report.json"
            quality_report = {
                "station_id": station_id,
                "overall_completeness": quality_metrics.overall_completeness,
                "quality_score": quality_metrics.quality_score,
                "recommendation": quality_metrics.recommendation,
                "data_years": len(quality_metrics.annual_completeness),
                "wet_day_frequency": quality_metrics.wet_day_frequency,
                "suspect_data_percentage": quality_metrics.suspect_data_percentage,
                "significant_gaps": quality_metrics.significant_gaps,
                "quality_notes": quality_metrics.quality_notes,
                "analysis_date": pd.Timestamp.now().isoformat()
            }
            
            with open(quality_report_path, 'w') as f:
                json.dump(quality_report, f, indent=2, default=str)
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            # Verify quality report was generated
            report_exists = quality_report_path.exists()
            report_size = quality_report_path.stat().st_size if report_exists else 0
            
            # Verify quality metrics are reasonable
            metrics_valid = (
                0.0 <= quality_metrics.overall_completeness <= 1.0 and
                0.0 <= quality_metrics.quality_score <= 1.0 and
                quality_metrics.wet_day_frequency >= 0.0 and
                0.0 <= quality_metrics.suspect_data_percentage <= 100.0
            )
            
            success = report_exists and report_size > 0 and metrics_valid
            
            self.logger.log_test_result(
                f"quality_report_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "report_path": str(quality_report_path),
                    "report_exists": report_exists,
                    "report_size_bytes": report_size,
                    "metrics_valid": metrics_valid,
                    "overall_completeness": quality_metrics.overall_completeness,
                    "quality_score": quality_metrics.quality_score,
                    "recommendation": quality_metrics.recommendation
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"quality_report_{station_id}")
            return False
    
    def test_station_metadata_extraction(self, station_id: str) -> bool:
        """
        Test extracting and storing station metadata.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if metadata extracted and stored successfully, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Get station information
            station_info_dict = self.downloader.get_station_info(station_id)
            
            if not station_info_dict:
                self.logger.log_test_result(
                    f"metadata_extraction_{station_id}",
                    False,
                    0.0,
                    details={"station_id": station_id, "error": "Station not found"}
                )
                return False
            
            # Create StationInfo object
            station_info = StationInfo(
                id=station_info_dict['id'],
                name=station_info_dict['name'],
                latitude=station_info_dict['latitude'],
                longitude=station_info_dict['longitude'],
                elevation_m=station_info_dict.get('elevation_m'),
                data_years=station_info_dict.get('years_available', 0),
                first_year=station_info_dict.get('first_year', 0),
                last_year=station_info_dict.get('last_year', 0)
            )
            
            # Create station directory for organized output
            station_dir = self.directory_manager.create_station_directory(station_id)
            data_dir = station_dir / "data"
            
            # Save station metadata
            metadata_path = data_dir / "station_info.json"
            metadata = {
                "id": station_info.id,
                "name": station_info.name,
                "latitude": station_info.latitude,
                "longitude": station_info.longitude,
                "elevation_m": station_info.elevation_m,
                "data_years": station_info.data_years,
                "first_year": station_info.first_year,
                "last_year": station_info.last_year,
                "coordinates": station_info.coordinates,
                "extraction_date": pd.Timestamp.now().isoformat()
            }
            
            with open(metadata_path, 'w') as f:
                json.dump(metadata, f, indent=2, default=str)
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            # Verify metadata was saved
            metadata_exists = metadata_path.exists()
            metadata_size = metadata_path.stat().st_size if metadata_exists else 0
            
            # Verify metadata content is valid
            metadata_valid = (
                station_info.id == station_id and
                -90 <= station_info.latitude <= 90 and
                -180 <= station_info.longitude <= 180 and
                station_info.data_years >= 0
            )
            
            success = metadata_exists and metadata_size > 0 and metadata_valid
            
            self.logger.log_test_result(
                f"metadata_extraction_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "metadata_path": str(metadata_path),
                    "metadata_exists": metadata_exists,
                    "metadata_size_bytes": metadata_size,
                    "metadata_valid": metadata_valid,
                    "station_name": station_info.name,
                    "coordinates": station_info.coordinates,
                    "data_years": station_info.data_years
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"metadata_extraction_{station_id}")
            return False
    
    def test_quality_threshold_validation(self, station_id: str, 
                                        min_completeness: float = 0.7,
                                        min_quality_score: float = 0.6) -> bool:
        """
        Test quality threshold validation.
        
        Args:
            station_id: GHCN station identifier
            min_completeness: Minimum data completeness threshold
            min_quality_score: Minimum quality score threshold
            
        Returns:
            True if quality validation works correctly, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Download and analyze data
            dly_file_path = self.downloader.download_station_data(station_id)
            quality_metrics = self.quality_analyzer.analyze_station_quality(station_id, dly_file_path)
            
            # Apply quality thresholds
            completeness_pass = quality_metrics.overall_completeness >= min_completeness
            quality_pass = quality_metrics.quality_score >= min_quality_score
            overall_pass = completeness_pass and quality_pass
            
            # Determine recommendation based on thresholds
            if overall_pass:
                expected_recommendation = "good" if quality_metrics.quality_score >= 0.8 else "acceptable"
            else:
                expected_recommendation = "poor"
            
            # Check if recommendation aligns with thresholds
            recommendation_appropriate = (
                (overall_pass and quality_metrics.recommendation in ["excellent", "good", "acceptable"]) or
                (not overall_pass and quality_metrics.recommendation in ["poor", "unacceptable"])
            )
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            # Validation is successful if thresholds work as expected
            success = True  # The validation logic itself should work
            
            self.logger.log_test_result(
                f"quality_validation_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "min_completeness": min_completeness,
                    "min_quality_score": min_quality_score,
                    "actual_completeness": quality_metrics.overall_completeness,
                    "actual_quality_score": quality_metrics.quality_score,
                    "completeness_pass": completeness_pass,
                    "quality_pass": quality_pass,
                    "overall_pass": overall_pass,
                    "recommendation": quality_metrics.recommendation,
                    "recommendation_appropriate": recommendation_appropriate
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"quality_validation_{station_id}")
            return False
    
    def test_data_gap_identification(self, station_id: str) -> bool:
        """
        Test identification of significant data gaps.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if gap identification works correctly, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Download and analyze data
            dly_file_path = self.downloader.download_station_data(station_id)
            quality_metrics = self.quality_analyzer.analyze_station_quality(station_id, dly_file_path)
            
            # Check gap identification
            gaps_identified = len(quality_metrics.significant_gaps) >= 0  # Should always be non-negative
            
            # Verify gap format if gaps exist
            gaps_valid = True
            for gap in quality_metrics.significant_gaps:
                if not (isinstance(gap, (list, tuple)) and len(gap) == 2):
                    gaps_valid = False
                    break
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            success = gaps_identified and gaps_valid
            
            self.logger.log_test_result(
                f"gap_identification_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "gaps_count": len(quality_metrics.significant_gaps),
                    "gaps_identified": gaps_identified,
                    "gaps_valid": gaps_valid,
                    "sample_gaps": quality_metrics.significant_gaps[:3]  # First 3 gaps
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"gap_identification_{station_id}")
            return False
    
    def test_quality_notes_generation(self, station_id: str) -> bool:
        """
        Test generation of quality notes and recommendations.
        
        Args:
            station_id: GHCN station identifier
            
        Returns:
            True if quality notes are generated appropriately, False otherwise
        """
        try:
            start_time = pd.Timestamp.now()
            
            # Download and analyze data
            dly_file_path = self.downloader.download_station_data(station_id)
            quality_metrics = self.quality_analyzer.analyze_station_quality(station_id, dly_file_path)
            
            # Check quality notes
            notes_exist = len(quality_metrics.quality_notes) >= 0
            notes_are_strings = all(isinstance(note, str) for note in quality_metrics.quality_notes)
            
            # Check recommendation is valid
            valid_recommendations = ["excellent", "good", "acceptable", "poor", "unacceptable", "unknown"]
            recommendation_valid = quality_metrics.recommendation in valid_recommendations
            
            execution_time = (pd.Timestamp.now() - start_time).total_seconds()
            
            success = notes_exist and notes_are_strings and recommendation_valid
            
            self.logger.log_test_result(
                f"quality_notes_{station_id}",
                success,
                execution_time,
                details={
                    "station_id": station_id,
                    "notes_count": len(quality_metrics.quality_notes),
                    "notes_exist": notes_exist,
                    "notes_are_strings": notes_are_strings,
                    "recommendation": quality_metrics.recommendation,
                    "recommendation_valid": recommendation_valid,
                    "sample_notes": quality_metrics.quality_notes[:3]  # First 3 notes
                }
            )
            
            return success
            
        except Exception as e:
            self.logger.log_error_with_context(e, f"quality_notes_{station_id}")
            return False
    
    def test_multiple_station_quality_assessment(self, station_ids: List[str]) -> Dict[str, bool]:
        """
        Test quality assessment 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 quality assessment: {len(station_ids)} stations")
        
        try:
            for i, station_id in enumerate(station_ids):
                try:
                    # Test all quality assessment components for each station
                    report_success = self.test_quality_report_generation(station_id)
                    metadata_success = self.test_station_metadata_extraction(station_id)
                    validation_success = self.test_quality_threshold_validation(station_id)
                    gaps_success = self.test_data_gap_identification(station_id)
                    notes_success = self.test_quality_notes_generation(station_id)
                    
                    # Overall success for this station
                    station_success = all([
                        report_success, metadata_success, validation_success,
                        gaps_success, notes_success
                    ])
                    results[station_id] = station_success
                    
                    self.logger.log_test_result(
                        f"multi_quality_{station_id}",
                        station_success,
                        0.0,
                        details={
                            "station_id": station_id,
                            "progress": f"{i+1}/{len(station_ids)}",
                            "report": report_success,
                            "metadata": metadata_success,
                            "validation": validation_success,
                            "gaps": gaps_success,
                            "notes": notes_success
                        }
                    )
                    
                except Exception as e:
                    results[station_id] = False
                    self.logger.log_error_with_context(e, f"multi_quality_{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 quality assessments completed: {success_rate:.1%} success rate"
            )
            
            return results
            
        except Exception as e:
            self.logger.log_error_with_context(e, "multiple_station_quality_assessment")
            self.logger.end_test_execution(False, f"Multiple quality assessments failed: {e}")
            return results
    
    def run_comprehensive_tests(self) -> Dict[str, bool]:
        """
        Run comprehensive quality assessment tests.
        
        Returns:
            Dictionary of test results
        """
        results = {}
        
        self.logger.start_test_execution("Comprehensive Quality Assessment Tests")
        
        try:
            # Test with Salt Lake City station (gold standard)
            slc_station = "USC00426128"
            
            # Individual component tests
            results["quality_report"] = self.test_quality_report_generation(slc_station)
            results["metadata_extraction"] = self.test_station_metadata_extraction(slc_station)
            results["quality_validation"] = self.test_quality_threshold_validation(slc_station)
            results["gap_identification"] = self.test_data_gap_identification(slc_station)
            results["quality_notes"] = self.test_quality_notes_generation(slc_station)
            
            # Multiple station test
            test_stations = ["USC00426128", "USW00024127", "USC00305426"]
            multi_results = self.test_multiple_station_quality_assessment(test_stations)
            results["multiple_assessments"] = 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 TestQualityAssessmentWorkflow:
    """Test suite for quality assessment workflow components."""
    
    def setup_method(self):
        """Set up test environment."""
        self.temp_dir = tempfile.mkdtemp(prefix="precipgen_quality_test_")
        self.tester = QualityAssessmentTester(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_quality_assessment_tester_initialization(self):
        """Test that QualityAssessmentTester 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.quality_analyzer is not None
        assert self.tester.directory_manager is not None
    
    def test_quality_report_generation_salt_lake_city(self):
        """Test generating quality report for Salt Lake City (gold standard)."""
        slc_station = "USC00426128"
        
        success = self.tester.test_quality_report_generation(slc_station)
        assert success == True, "Should successfully generate quality report"
        
        # Verify quality report exists in organized directory
        station_dir = Path(self.temp_dir) / slc_station
        quality_report_path = station_dir / "data" / "quality_report.json"
        
        assert quality_report_path.exists(), "Quality report should exist in data subdirectory"
        assert quality_report_path.stat().st_size > 0, "Quality report should have content"
        
        # Verify report content
        with open(quality_report_path, 'r') as f:
            report_data = json.load(f)
        
        assert report_data["station_id"] == slc_station, "Report should contain correct station ID"
        assert "overall_completeness" in report_data, "Report should contain completeness metric"
        assert "quality_score" in report_data, "Report should contain quality score"
    
    def test_station_metadata_extraction_salt_lake_city(self):
        """Test extracting station metadata for Salt Lake City."""
        slc_station = "USC00426128"
        
        success = self.tester.test_station_metadata_extraction(slc_station)
        assert success == True, "Should successfully extract station metadata"
        
        # Verify metadata exists in organized directory
        station_dir = Path(self.temp_dir) / slc_station
        metadata_path = station_dir / "data" / "station_info.json"
        
        assert metadata_path.exists(), "Station metadata should exist in data subdirectory"
        assert metadata_path.stat().st_size > 0, "Metadata should have content"
        
        # Verify metadata content
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        assert metadata["id"] == slc_station, "Metadata should contain correct station ID"
        assert "latitude" in metadata, "Metadata should contain latitude"
        assert "longitude" in metadata, "Metadata should contain longitude"
        assert "name" in metadata, "Metadata should contain station name"
    
    def test_quality_threshold_validation(self):
        """Test quality threshold validation."""
        slc_station = "USC00426128"
        
        success = self.tester.test_quality_threshold_validation(slc_station)
        assert success == True, "Quality threshold validation should work correctly"
    
    def test_data_gap_identification(self):
        """Test data gap identification."""
        slc_station = "USC00426128"
        
        success = self.tester.test_data_gap_identification(slc_station)
        assert success == True, "Data gap identification should work correctly"
    
    def test_quality_notes_generation(self):
        """Test quality notes generation."""
        slc_station = "USC00426128"
        
        success = self.tester.test_quality_notes_generation(slc_station)
        assert success == True, "Quality notes generation should work correctly"
    
    def test_comprehensive_workflow(self):
        """Test the complete quality assessment 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 data-dependent variations)
        success_rate = sum(results.values()) / len(results)
        assert success_rate >= 0.7, f"Success rate too low: {success_rate:.2f}"
