# Module 17: Comprehensive Integration Tests
# tests/test_integration.py

import pytest
import asyncio
import httpx
import time
from typing import List, Dict, Any
import statistics
import json
import os
from datetime import datetime

# Test configuration
API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:8000")
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
AZURE_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")

class TestSemanticSearch:
    """Test semantic search functionality (Exercise 1)."""
    
    @pytest.fixture
    async def client(self):
        async with httpx.AsyncClient(base_url=f"{API_BASE_URL}:8001") as client:
            yield client
    
    @pytest.mark.asyncio
    async def test_embedding_generation(self, client):
        """Test embedding generation endpoint."""
        response = await client.post(
            "/embeddings",
            json={
                "text": "Advanced AI security camera with facial recognition",
                "model": "text-embedding-ada-002"
            }
        )
        
        assert response.status_code == 200
        data = response.json()
        
        # Validate response structure
        assert "embedding" in data
        assert "model" in data
        assert "usage" in data
        
        # Validate embedding
        embedding = data["embedding"]
        assert isinstance(embedding, list)
        assert len(embedding) == 1536  # Ada-002 dimensions
        assert all(isinstance(x, float) for x in embedding)
    
    @pytest.mark.asyncio
    async def test_semantic_search(self, client):
        """Test semantic search functionality."""
        response = await client.post(
            "/search",
            json={
                "query": "home security camera with night vision",
                "top_k": 5
            }
        )
        
        assert response.status_code == 200
        data = response.json()
        
        # Validate response
        assert "results" in data
        assert "query" in data
        assert "search_time_ms" in data
        
        # Check results
        results = data["results"]
        assert len(results) <= 5
        
        # Verify relevance (first result should be most relevant)
        if results:
            first_result = results[0]
            assert first_result["score"] >= 0.7  # Reasonable threshold
            assert "product" in first_result
            assert "explanation" in first_result
    
    @pytest.mark.asyncio
    async def test_model_comparison(self, client):
        """Test different embedding models."""
        models = ["text-embedding-ada-002", "text-embedding-3-small"]
        query = "quantum computing fundamentals"
        
        results = {}
        for model in models:
            # Skip if model not available
            try:
                response = await client.post(
                    "/embeddings",
                    json={"text": query, "model": model}
                )
                if response.status_code == 200:
                    results[model] = response.json()
            except:
                continue
        
        # Compare if we have multiple models
        if len(results) > 1:
            # Embeddings should be different but similar in magnitude
            for model, data in results.items():
                embedding = data["embedding"]
                magnitude = sum(x**2 for x in embedding) ** 0.5
                assert 0.9 < magnitude < 1.1  # Normalized embeddings

class TestRAGSystem:
    """Test RAG system functionality (Exercise 2)."""
    
    @pytest.fixture
    async def client(self):
        async with httpx.AsyncClient(base_url=f"{API_BASE_URL}:8002") as client:
            yield client
    
    @pytest.mark.asyncio
    async def test_rag_query(self, client):
        """Test RAG query processing."""
        response = await client.post(
            "/query",
            json={
                "query": "What are best practices for AI model deployment?",
                "mode": "detailed",
                "include_sources": True
            }
        )
        
        assert response.status_code == 200
        data = response.json()
        
        # Validate response structure
        assert "response" in data
        assert "sources" in data
        assert "confidence" in data
        assert "metadata" in data
        
        # Check response quality
        assert len(data["response"]) > 100  # Non-trivial response
        assert data["confidence"] > 0.5  # Reasonable confidence
        
        # Verify sources if included
        if data["sources"]:
            source = data["sources"][0]
            assert "document" in source
            assert "relevance_score" in source
    
    @pytest.mark.asyncio
    async def test_rag_streaming(self, client):
        """Test RAG streaming response."""
        chunks = []
        
        async with client.stream(
            "POST",
            "/query/stream",
            json={"query": "Explain feature engineering", "stream": True}
        ) as response:
            async for chunk in response.aiter_text():
                chunks.append(chunk)
        
        # Verify streaming worked
        assert len(chunks) > 0
        complete_response = "".join(chunks)
        assert len(complete_response) > 50
    
    @pytest.mark.asyncio
    async def test_document_indexing(self, client):
        """Test document indexing pipeline."""
        # Skip if no test documents
        test_doc = "data/test_document.txt"
        if not os.path.exists(test_doc):
            pytest.skip("Test document not found")
        
        response = await client.post(
            "/index",
            json={"documents": [test_doc]}
        )
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["documents_processed"] > 0
        assert data["chunks_created"] > 0
        assert len(data["errors"]) == 0

class TestMultiModelOrchestration:
    """Test multi-model orchestration (Exercise 3)."""
    
    @pytest.fixture
    async def client(self):
        async with httpx.AsyncClient(
            base_url=API_BASE_URL,
            timeout=30.0
        ) as client:
            yield client
    
    @pytest.mark.asyncio
    async def test_model_routing(self, client):
        """Test intelligent model routing."""
        test_cases = [
            {
                "prompt": "Write a Python function to sort a list",
                "task_type": "code-generation",
                "expected_models": ["code-davinci-002", "gpt-4"]
            },
            {
                "prompt": "Summarize this article about climate change",
                "task_type": "summarization",
                "expected_models": ["gpt-3.5-turbo", "gpt-4"]
            },
            {
                "prompt": "Hi, how are you?",
                "task_type": "conversation",
                "expected_models": ["gpt-3.5-turbo", "llama-2-70b"]
            }
        ]
        
        for test_case in test_cases:
            response = await client.post(
                "/v1/completions",
                json={
                    "prompt": test_case["prompt"],
                    "task_type": test_case["task_type"]
                }
            )
            
            assert response.status_code == 200
            data = response.json()
            
            # Verify routing
            assert data["model_used"] in test_case["expected_models"]
            assert "response" in data
            assert "cost" in data
            assert data["cost"] > 0
    
    @pytest.mark.asyncio
    async def test_fallback_mechanism(self, client):
        """Test model fallback on failure."""
        # Force primary model failure
        response = await client.post(
            "/v1/completions",
            json={
                "prompt": "Test prompt",
                "model_preferences": {
                    "exclude_providers": ["azure_openai"]  # Force fallback
                }
            }
        )
        
        assert response.status_code == 200
        data = response.json()
        
        # Should use fallback model
        assert data["model_used"] != "gpt-4"  # Primary model
        assert "metadata" in data
        assert data["metadata"].get("fallback_used", False)
    
    @pytest.mark.asyncio
    async def test_cost_optimization(self, client):
        """Test cost optimization features."""
        # Request with budget constraint
        response = await client.post(
            "/v1/completions",
            json={
                "prompt": "Explain quantum computing in simple terms",
                "budget_limit": 0.01,  # $0.01 limit
                "max_tokens": 100
            }
        )
        
        assert response.status_code == 200
        data = response.json()
        
        # Verify budget constraint respected
        assert data["cost"] <= 0.01
        # Should use cheaper model
        assert data["model_used"] in ["gpt-3.5-turbo", "llama-2-70b"]
    
    @pytest.mark.asyncio
    async def test_concurrent_requests(self, client):
        """Test system under concurrent load."""
        num_requests = 20
        
        async def make_request(i):
            try:
                response = await client.post(
                    "/v1/completions",
                    json={
                        "prompt": f"Test request {i}",
                        "max_tokens": 50
                    }
                )
                return {
                    "success": response.status_code == 200,
                    "latency": response.elapsed.total_seconds()
                }
            except Exception as e:
                return {"success": False, "error": str(e)}
        
        # Run concurrent requests
        tasks = [make_request(i) for i in range(num_requests)]
        results = await asyncio.gather(*tasks)
        
        # Analyze results
        successful = sum(1 for r in results if r.get("success", False))
        success_rate = successful / num_requests
        
        latencies = [r["latency"] for r in results if "latency" in r]
        avg_latency = statistics.mean(latencies) if latencies else 0
        p95_latency = statistics.quantiles(latencies, n=20)[18] if latencies else 0
        
        # Assertions
        assert success_rate >= 0.95  # 95% success rate
        assert avg_latency < 2.0  # Average under 2 seconds
        assert p95_latency < 5.0  # P95 under 5 seconds

class TestEndToEnd:
    """End-to-end integration tests."""
    
    @pytest.mark.asyncio
    async def test_complete_workflow(self):
        """Test complete AI integration workflow."""
        async with httpx.AsyncClient(timeout=30.0) as client:
            # 1. Generate embedding
            embedding_response = await client.post(
                f"{API_BASE_URL}:8001/embeddings",
                json={"text": "AI-powered code analysis tool"}
            )
            assert embedding_response.status_code == 200
            
            # 2. Search using embedding
            search_response = await client.post(
                f"{API_BASE_URL}:8001/search",
                json={"query": "code analysis AI tools", "top_k": 3}
            )
            assert search_response.status_code == 200
            
            # 3. Use search results in RAG
            search_results = search_response.json()["results"]
            if search_results:
                context = search_results[0]["product"]["description"]
                
                rag_response = await client.post(
                    f"{API_BASE_URL}:8002/query",
                    json={
                        "query": f"Based on this product: {context}, what are similar tools?",
                        "mode": "concise"
                    }
                )
                assert rag_response.status_code == 200
            
            # 4. Use orchestrator for optimization
            orchestrator_response = await client.post(
                f"{API_BASE_URL}/v1/completions",
                json={
                    "prompt": "Compare different code analysis tools",
                    "task_type": "analysis",
                    "budget_limit": 0.05
                }
            )
            assert orchestrator_response.status_code == 200

@pytest.mark.benchmark
class TestPerformance:
    """Performance benchmarking tests."""
    
    @pytest.mark.asyncio
    async def test_embedding_performance(self):
        """Benchmark embedding generation performance."""
        async with httpx.AsyncClient(base_url=f"{API_BASE_URL}:8001") as client:
            texts = [
                f"Test text {i} for embedding generation benchmark"
                for i in range(100)
            ]
            
            start_time = time.time()
            
            # Test batch processing
            tasks = []
            for text in texts:
                task = client.post(
                    "/embeddings",
                    json={"text": text}
                )
                tasks.append(task)
            
            responses = await asyncio.gather(*tasks, return_exceptions=True)
            
            end_time = time.time()
            duration = end_time - start_time
            
            # Calculate metrics
            successful = sum(
                1 for r in responses 
                if not isinstance(r, Exception) and r.status_code == 200
            )
            
            print(f"\nEmbedding Performance:")
            print(f"Total requests: {len(texts)}")
            print(f"Successful: {successful}")
            print(f"Total time: {duration:.2f}s")
            print(f"Throughput: {successful/duration:.2f} req/s")
            print(f"Avg latency: {duration/len(texts)*1000:.2f}ms")
            
            # Performance assertions
            assert successful/len(texts) > 0.95  # 95% success rate
            assert duration < 60  # Complete within 1 minute

def pytest_sessionfinish(session, exitstatus):
    """Generate test report after all tests."""
    print("\n" + "="*50)
    print("Module 17 Integration Test Summary")
    print("="*50)
    
    # Summary would be generated here
    print(f"Exit Status: {exitstatus}")
    print(f"Tests Run: {session.testscollected}")
    print(f"Timestamp: {datetime.now().isoformat()}")

if __name__ == "__main__":
    # Run tests with coverage
    pytest.main([
        __file__,
        "-v",
        "--tb=short",
        "--cov=.",
        "--cov-report=html",
        "--cov-report=term-missing"
    ])