"""
Unit tests for RAG pipeline core functionality
"""

import pytest
import time
from unittest.mock import Mock, patch

from grass_rag.core.pipeline import OptimizedRAGPipeline
from grass_rag.core.models import RAGConfig, RAGResponse


class TestOptimizedRAGPipeline:
    """Test cases for OptimizedRAGPipeline"""
    
    def setup_method(self):
        """Set up test environment"""
        self.config = RAGConfig(
            cache_size=100,
            max_response_time=5.0,
            template_threshold=0.8
        )
        self.pipeline = OptimizedRAGPipeline(self.config.to_dict())
    
    def test_initialization(self):
        """Test pipeline initialization"""
        assert self.pipeline.config.cache_size == 100
        assert self.pipeline.config.max_response_time == 5.0
        assert self.pipeline.templates is not None
        assert self.pipeline.cache is not None
        assert self.pipeline.stats["total_queries"] == 0
    
    def test_template_query(self):
        """Test query that should match template"""
        query = "calculate slope from DEM"
        
        answer, sources, metrics = self.pipeline.query(query)
        
        assert answer is not None
        assert len(answer) > 0
        assert isinstance(sources, list)
        assert isinstance(metrics, dict)
        assert metrics["method"] == "template"
        assert metrics["quality_score"] >= 0.9
        assert metrics["total_time"] < 1.0  # Should be fast
    
    def test_cache_functionality(self):
        """Test cache hit on repeated query"""
        query = "import raster data"
        
        # First query
        answer1, sources1, metrics1 = self.pipeline.query(query)
        
        # Second query (should hit cache)
        answer2, sources2, metrics2 = self.pipeline.query(query)
        
        assert answer1 == answer2
        assert metrics2["total_time"] < metrics1["total_time"]  # Cache should be faster
    
    def test_fallback_query(self):
        """Test query that should use fallback"""
        query = "some very specific unknown operation xyz123"
        
        answer, sources, metrics = self.pipeline.query(query)
        
        assert answer is not None
        assert len(answer) > 0
        assert metrics["method"] == "enhanced_fallback"
        assert metrics["quality_score"] >= 0.8  # Fallback should still be decent quality
    
    def test_batch_query(self):
        """Test batch query processing"""
        queries = [
            "calculate slope",
            "import raster",
            "create buffer"
        ]
        
        results = self.pipeline.batch_query(queries)
        
        assert len(results) == len(queries)
        for answer, sources, metrics in results:
            assert answer is not None
            assert isinstance(sources, list)
            assert isinstance(metrics, dict)
    
    def test_performance_requirements(self):
        """Test that performance requirements are met"""
        test_queries = [
            "calculate slope from DEM",
            "import raster file",
            "export vector data",
            "create contour lines",
            "vector overlay analysis",
            "watershed analysis",
            "create buffer zones",
            "raster calculations",
            "classify raster data",
            "interpolate surface"
        ]
        
        quality_scores = []
        response_times = []
        
        for query in test_queries:
            start_time = time.time()
            answer, sources, metrics = self.pipeline.query(query)
            response_time = time.time() - start_time
            
            quality_scores.append(metrics["quality_score"])
            response_times.append(response_time)
        
        # Check performance requirements
        avg_quality = sum(quality_scores) / len(quality_scores)
        avg_response_time = sum(response_times) / len(response_times)
        
        assert avg_quality >= 0.9, f"Average quality {avg_quality} below 0.9 requirement"
        assert avg_response_time < 5.0, f"Average response time {avg_response_time}s above 5s requirement"
        
        # Check that at least 90% of queries meet quality requirement
        high_quality_count = sum(1 for q in quality_scores if q >= 0.9)
        quality_percentage = (high_quality_count / len(quality_scores)) * 100
        assert quality_percentage >= 90, f"Only {quality_percentage}% of queries meet quality requirement"
    
    def test_configuration_update(self):
        """Test configuration updates"""
        original_cache_size = self.pipeline.config.cache_size
        
        self.pipeline.configure(cache_size=500)
        
        assert self.pipeline.config.cache_size == 500
        assert self.pipeline.config.cache_size != original_cache_size
    
    def test_performance_report(self):
        """Test performance report generation"""
        # Generate some queries first
        queries = ["calculate slope", "import data", "export raster"]
        for query in queries:
            self.pipeline.query(query)
        
        report = self.pipeline.get_performance_report()
        
        assert "performance_summary" in report
        assert "quality_analysis" in report
        assert "speed_analysis" in report
        assert "target_achievement" in report
        
        summary = report["performance_summary"]
        assert summary["total_queries"] == len(queries)
        assert "avg_quality_score" in summary
        assert "avg_response_time" in summary
    
    def test_error_handling(self):
        """Test error handling for invalid inputs"""
        # Test empty query
        with pytest.raises(Exception):
            self.pipeline.query("")
        
        # Test None query
        with pytest.raises(Exception):
            self.pipeline.query(None)
    
    def test_stats_tracking(self):
        """Test statistics tracking"""
        initial_queries = self.pipeline.stats["total_queries"]
        
        self.pipeline.query("test query")
        
        assert self.pipeline.stats["total_queries"] == initial_queries + 1
        assert len(self.pipeline.stats["quality_scores"]) == initial_queries + 1
        assert len(self.pipeline.stats["response_times"]) == initial_queries + 1
    
    @patch('grass_rag.core.pipeline.ModelDownloadManager')
    def test_model_loading(self, mock_model_manager):
        """Test model loading functionality"""
        mock_manager = Mock()
        mock_manager.verify_models.return_value = True
        mock_model_manager.return_value = mock_manager
        
        pipeline = OptimizedRAGPipeline()
        pipeline._ensure_models_loaded()
        
        mock_manager.verify_models.assert_called_once()
    
    def test_cache_stats(self):
        """Test cache statistics"""
        # Generate some queries to populate cache
        queries = ["slope calculation", "data import", "raster export"]
        for query in queries:
            self.pipeline.query(query)
        
        cache_stats = self.pipeline.get_cache_stats()
        
        assert "l1_size" in cache_stats
        assert "l2_size" in cache_stats
        assert "hit_rate" in cache_stats
        assert isinstance(cache_stats["hit_rate"], (int, float))
    
    def test_template_stats(self):
        """Test template statistics"""
        template_stats = self.pipeline.get_template_stats()
        
        assert "total_templates" in template_stats
        assert "categories" in template_stats
        assert "avg_quality_score" in template_stats
        assert template_stats["total_templates"] > 0
        assert template_stats["avg_quality_score"] >= 0.9


if __name__ == "__main__":
    pytest.main([__file__])