#!/usr/bin/env python3
"""
Performance tests for user authentication feature

This test file focuses on performance testing of user_id queries, indexes,
and analytics functions to ensure the user authentication feature doesn't
negatively impact system performance.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Load environment variables
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # dotenv not available, continue without it

import asyncio
import time
import random
import string
from datetime import datetime, timedelta
from typing import List, Dict, Any

from analytics.collectors import create_session, track_message, track_search
from analytics.dashboard import (
    get_user_statistics,
    get_user_analytics,
    get_authenticated_vs_anonymous_metrics,
    get_basic_stats,
    get_hourly_message_stats
)
from analytics.database import (
    get_sessions_collection,
    get_messages_collection,
    get_search_analytics_collection
)


class TestDatabaseIndexPerformance:
    """Test performance of database indexes for user_id queries"""
    
    async def test_user_id_index_performance(self):
        """Test performance of user_id index queries"""
        print("🔍 Testing user_id index performance...")
        
        # Create test data with multiple users
        user_ids = [f"perf_user_{i}" for i in range(20)]
        sessions_per_user = 3
        messages_per_session = 5
        
        print(f"Creating test data: {len(user_ids)} users, {sessions_per_user} sessions each, {messages_per_session} messages each")
        
        # Create test data
        start_time = time.time()
        all_sessions = []
        
        for user_id in user_ids:
            for session_num in range(sessions_per_user):
                session = await create_session(user_id=user_id)
                all_sessions.append((session, user_id))
                
                for msg_num in range(messages_per_session):
                    await track_message(
                        session_id=session.session_id,
                        prompt_length=random.randint(20, 100),
                        response_length=random.randint(50, 200),
                        response_time_ms=random.randint(500, 2000),
                        used_search=random.choice([True, False]),
                        user_id=user_id
                    )
        
        data_creation_time = time.time() - start_time
        print(f"✅ Test data created in {data_creation_time:.2f} seconds")
        
        # Test query performance
        sessions_collection = await get_sessions_collection()
        messages_collection = await get_messages_collection()
        
        if sessions_collection and messages_collection:
            # Test individual user queries
            print("Testing individual user queries...")
            start_time = time.time()
            
            for user_id in user_ids:
                user_sessions = await sessions_collection.count_documents({"user_id": user_id})
                assert user_sessions == sessions_per_user
                
                user_messages = await messages_collection.count_documents({"user_id": user_id})
                assert user_messages == sessions_per_user * messages_per_session
            
            individual_query_time = time.time() - start_time
            avg_query_time = individual_query_time / len(user_ids)
            
            print(f"✅ Individual user queries: {individual_query_time:.2f}s total, {avg_query_time:.4f}s average")
            
            # Performance assertion
            assert avg_query_time < 0.1, f"Individual queries too slow: {avg_query_time:.4f}s average"
            
            # Test bulk queries
            print("Testing bulk user queries...")
            start_time = time.time()
            
            # Query all authenticated sessions
            auth_sessions = await sessions_collection.count_documents({"user_id": {"$ne": None}})
            assert auth_sessions == len(user_ids) * sessions_per_user
            
            # Query all authenticated messages
            auth_messages = await messages_collection.count_documents({"user_id": {"$ne": None}})
            assert auth_messages == len(user_ids) * sessions_per_user * messages_per_session
            
            bulk_query_time = time.time() - start_time
            print(f"✅ Bulk queries: {bulk_query_time:.2f}s")
            
            # Performance assertion
            assert bulk_query_time < 2.0, f"Bulk queries too slow: {bulk_query_time:.2f}s"
            
        else:
            print("⚠️  Database not available - skipping index performance test")
    
    async def test_compound_index_performance(self):
        """Test performance of compound (user_id, timestamp) index queries"""
        print("🔍 Testing compound index performance...")
        
        # Create test data with timestamps spread over time
        user_id = "compound_perf_user"
        session = await create_session(user_id=user_id)
        
        num_messages = 100
        print(f"Creating {num_messages} messages with varied timestamps...")
        
        start_time = time.time()
        base_time = datetime.utcnow()
        
        for i in range(num_messages):
            # Create messages with timestamps spread over the last 24 hours
            timestamp_offset = timedelta(hours=random.uniform(0, 24))
            message_time = base_time - timestamp_offset
            
            await track_message(
                session_id=session.session_id,
                prompt_length=50,
                response_length=100,
                response_time_ms=1000,
                user_id=user_id
            )
            
            # Small delay to ensure different timestamps
            await asyncio.sleep(0.001)
        
        data_creation_time = time.time() - start_time
        print(f"✅ Test data created in {data_creation_time:.2f} seconds")
        
        # Test compound queries
        messages_collection = await get_messages_collection()
        
        if messages_collection:
            # Test various time range queries
            time_ranges = [
                ("1 hour", timedelta(hours=1)),
                ("6 hours", timedelta(hours=6)),
                ("12 hours", timedelta(hours=12)),
                ("24 hours", timedelta(hours=24))
            ]
            
            for range_name, time_delta in time_ranges:
                start_time = time.time()
                cutoff_time = datetime.utcnow() - time_delta
                
                recent_messages = await messages_collection.count_documents({
                    "user_id": user_id,
                    "timestamp": {"$gte": cutoff_time}
                })
                
                query_time = time.time() - start_time
                print(f"✅ {range_name} query: {query_time:.4f}s ({recent_messages} messages)")
                
                # Performance assertion
                assert query_time < 0.5, f"{range_name} query too slow: {query_time:.4f}s"
        
        else:
            print("⚠️  Database not available - skipping compound index test")
    
    async def test_sparse_index_performance(self):
        """Test performance of sparse indexes with mixed null/non-null user_id values"""
        print("🔍 Testing sparse index performance...")
        
        # Create mixed data (authenticated and anonymous)
        num_auth_users = 10
        num_anon_sessions = 20
        messages_per_session = 5
        
        print(f"Creating mixed data: {num_auth_users} auth users, {num_anon_sessions} anon sessions")
        
        start_time = time.time()
        
        # Create authenticated user data
        auth_sessions = []
        for i in range(num_auth_users):
            user_id = f"sparse_user_{i}"
            session = await create_session(user_id=user_id)
            auth_sessions.append(session)
            
            for j in range(messages_per_session):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=50,
                    response_length=100,
                    response_time_ms=1000,
                    user_id=user_id
                )
        
        # Create anonymous user data
        anon_sessions = []
        for i in range(num_anon_sessions):
            session = await create_session(user_id=None)
            anon_sessions.append(session)
            
            for j in range(messages_per_session):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=50,
                    response_length=100,
                    response_time_ms=1000,
                    user_id=None
                )
        
        data_creation_time = time.time() - start_time
        print(f"✅ Mixed data created in {data_creation_time:.2f} seconds")
        
        # Test sparse index queries
        sessions_collection = await get_sessions_collection()
        messages_collection = await get_messages_collection()
        
        if sessions_collection and messages_collection:
            # Test authenticated user queries
            start_time = time.time()
            auth_session_count = await sessions_collection.count_documents({"user_id": {"$ne": None}})
            auth_query_time = time.time() - start_time
            
            assert auth_session_count >= num_auth_users
            print(f"✅ Authenticated sessions query: {auth_query_time:.4f}s ({auth_session_count} sessions)")
            
            # Test anonymous user queries
            start_time = time.time()
            anon_session_count = await sessions_collection.count_documents({"user_id": None})
            anon_query_time = time.time() - start_time
            
            assert anon_session_count >= num_anon_sessions
            print(f"✅ Anonymous sessions query: {anon_query_time:.4f}s ({anon_session_count} sessions)")
            
            # Test specific user queries
            start_time = time.time()
            specific_user_sessions = await sessions_collection.count_documents({"user_id": "sparse_user_0"})
            specific_query_time = time.time() - start_time
            
            assert specific_user_sessions == 1
            print(f"✅ Specific user query: {specific_query_time:.4f}s")
            
            # Performance assertions
            assert auth_query_time < 0.5, f"Auth query too slow: {auth_query_time:.4f}s"
            assert anon_query_time < 0.5, f"Anon query too slow: {anon_query_time:.4f}s"
            assert specific_query_time < 0.1, f"Specific query too slow: {specific_query_time:.4f}s"
        
        else:
            print("⚠️  Database not available - skipping sparse index test")


class TestAnalyticsFunctionPerformance:
    """Test performance of analytics functions with user authentication"""
    
    async def test_user_statistics_performance(self):
        """Test performance of get_user_statistics function"""
        print("📊 Testing user statistics performance...")
        
        # Create test data
        await self._create_performance_test_data()
        
        # Test get_user_statistics performance
        start_time = time.time()
        user_stats = await get_user_statistics()
        stats_time = time.time() - start_time
        
        assert isinstance(user_stats, dict)
        assert "unique_authenticated_users" in user_stats
        assert "authenticated_sessions" in user_stats
        assert "anonymous_sessions" in user_stats
        
        print(f"✅ User statistics query: {stats_time:.4f}s")
        
        # Performance assertion
        assert stats_time < 5.0, f"User statistics too slow: {stats_time:.4f}s"
    
    async def test_user_analytics_performance(self):
        """Test performance of get_user_analytics function"""
        print("📊 Testing individual user analytics performance...")
        
        # Create test user with substantial data
        user_id = "analytics_perf_user"
        session = await create_session(user_id=user_id)
        
        # Create many messages for this user
        num_messages = 50
        print(f"Creating {num_messages} messages for performance test...")
        
        for i in range(num_messages):
            await track_message(
                session_id=session.session_id,
                prompt_length=random.randint(20, 100),
                response_length=random.randint(50, 200),
                response_time_ms=random.randint(500, 2000),
                used_search=random.choice([True, False]),
                user_id=user_id
            )
        
        # Test get_user_analytics performance
        start_time = time.time()
        user_analytics = await get_user_analytics(user_id)
        analytics_time = time.time() - start_time
        
        assert isinstance(user_analytics, dict)
        assert user_analytics.get("user_id") == user_id
        assert user_analytics.get("total_messages") == num_messages
        
        print(f"✅ User analytics query: {analytics_time:.4f}s")
        
        # Performance assertion
        assert analytics_time < 3.0, f"User analytics too slow: {analytics_time:.4f}s"
    
    async def test_authenticated_vs_anonymous_performance(self):
        """Test performance of get_authenticated_vs_anonymous_metrics function"""
        print("📊 Testing authenticated vs anonymous metrics performance...")
        
        # Create mixed test data
        await self._create_performance_test_data()
        
        # Test get_authenticated_vs_anonymous_metrics performance
        start_time = time.time()
        comparison_metrics = await get_authenticated_vs_anonymous_metrics()
        comparison_time = time.time() - start_time
        
        assert isinstance(comparison_metrics, dict)
        assert "authenticated" in comparison_metrics
        assert "anonymous" in comparison_metrics
        assert "comparison" in comparison_metrics
        
        print(f"✅ Comparison metrics query: {comparison_time:.4f}s")
        
        # Performance assertion
        assert comparison_time < 5.0, f"Comparison metrics too slow: {comparison_time:.4f}s"
    
    async def test_basic_stats_with_filter_performance(self):
        """Test performance of get_basic_stats with user_id filter"""
        print("📊 Testing filtered basic stats performance...")
        
        # Create test data
        await self._create_performance_test_data()
        
        # Test unfiltered stats
        start_time = time.time()
        all_stats = await get_basic_stats()
        all_stats_time = time.time() - start_time
        
        print(f"✅ Unfiltered basic stats: {all_stats_time:.4f}s")
        
        # Test filtered stats
        start_time = time.time()
        filtered_stats = await get_basic_stats(user_id="perf_test_user_0")
        filtered_stats_time = time.time() - start_time
        
        assert "filtered_by_user_id" in filtered_stats
        print(f"✅ Filtered basic stats: {filtered_stats_time:.4f}s")
        
        # Performance assertions
        assert all_stats_time < 3.0, f"Unfiltered stats too slow: {all_stats_time:.4f}s"
        assert filtered_stats_time < 2.0, f"Filtered stats too slow: {filtered_stats_time:.4f}s"
    
    async def test_hourly_stats_with_filter_performance(self):
        """Test performance of get_hourly_message_stats with user_id filter"""
        print("📊 Testing filtered hourly stats performance...")
        
        # Create test data
        await self._create_performance_test_data()
        
        # Test unfiltered hourly stats
        start_time = time.time()
        all_hourly = await get_hourly_message_stats(hours=24)
        all_hourly_time = time.time() - start_time
        
        print(f"✅ Unfiltered hourly stats: {all_hourly_time:.4f}s")
        
        # Test filtered hourly stats
        start_time = time.time()
        filtered_hourly = await get_hourly_message_stats(hours=24, user_id="perf_test_user_0")
        filtered_hourly_time = time.time() - start_time
        
        print(f"✅ Filtered hourly stats: {filtered_hourly_time:.4f}s")
        
        # Performance assertions
        assert all_hourly_time < 3.0, f"Unfiltered hourly stats too slow: {all_hourly_time:.4f}s"
        assert filtered_hourly_time < 2.0, f"Filtered hourly stats too slow: {filtered_hourly_time:.4f}s"
    
    async def _create_performance_test_data(self):
        """Create test data for performance testing"""
        # Create authenticated users
        for i in range(5):
            user_id = f"perf_test_user_{i}"
            session = await create_session(user_id=user_id)
            
            # Create messages for each user
            for j in range(10):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=random.randint(20, 100),
                    response_length=random.randint(50, 200),
                    response_time_ms=random.randint(500, 2000),
                    used_search=random.choice([True, False]),
                    user_id=user_id
                )
        
        # Create anonymous users
        for i in range(3):
            session = await create_session(user_id=None)
            
            # Create messages for anonymous users
            for j in range(8):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=random.randint(20, 100),
                    response_length=random.randint(50, 200),
                    response_time_ms=random.randint(500, 2000),
                    used_search=random.choice([True, False]),
                    user_id=None
                )


class TestConcurrentUserPerformance:
    """Test performance with concurrent user operations"""
    
    async def test_concurrent_user_creation(self):
        """Test performance of concurrent user session creation"""
        print("🚀 Testing concurrent user creation performance...")
        
        num_concurrent_users = 20
        
        async def create_user_session(user_id: str):
            session = await create_session(user_id=user_id)
            
            # Create a few messages for each user
            for i in range(3):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=50,
                    response_length=100,
                    response_time_ms=1000,
                    user_id=user_id
                )
            
            return session
        
        # Create concurrent tasks
        start_time = time.time()
        tasks = [
            create_user_session(f"concurrent_user_{i}")
            for i in range(num_concurrent_users)
        ]
        
        sessions = await asyncio.gather(*tasks)
        concurrent_time = time.time() - start_time
        
        assert len(sessions) == num_concurrent_users
        print(f"✅ Concurrent user creation: {concurrent_time:.2f}s for {num_concurrent_users} users")
        
        # Performance assertion
        avg_time_per_user = concurrent_time / num_concurrent_users
        assert avg_time_per_user < 1.0, f"Concurrent creation too slow: {avg_time_per_user:.2f}s per user"
    
    async def test_concurrent_user_queries(self):
        """Test performance of concurrent user-specific queries"""
        print("🚀 Testing concurrent user query performance...")
        
        # Create test users first
        user_ids = [f"query_user_{i}" for i in range(10)]
        
        for user_id in user_ids:
            session = await create_session(user_id=user_id)
            await track_message(
                session_id=session.session_id,
                prompt_length=50,
                response_length=100,
                response_time_ms=1000,
                user_id=user_id
            )
        
        # Test concurrent queries
        async def query_user_analytics(user_id: str):
            return await get_user_analytics(user_id)
        
        start_time = time.time()
        tasks = [query_user_analytics(user_id) for user_id in user_ids]
        results = await asyncio.gather(*tasks)
        concurrent_query_time = time.time() - start_time
        
        assert len(results) == len(user_ids)
        for i, result in enumerate(results):
            assert result.get("user_id") == user_ids[i]
        
        print(f"✅ Concurrent user queries: {concurrent_query_time:.2f}s for {len(user_ids)} users")
        
        # Performance assertion
        avg_query_time = concurrent_query_time / len(user_ids)
        assert avg_query_time < 0.5, f"Concurrent queries too slow: {avg_query_time:.2f}s per query"


class TestMemoryPerformance:
    """Test memory usage with user authentication"""
    
    async def test_memory_usage_with_users(self):
        """Test that user_id fields don't significantly increase memory usage"""
        print("💾 Testing memory usage with user authentication...")
        
        import psutil
        import os
        
        # Get initial memory usage
        process = psutil.Process(os.getpid())
        initial_memory = process.memory_info().rss / 1024 / 1024  # MB
        
        # Create substantial amount of data
        num_users = 50
        messages_per_user = 20
        
        print(f"Creating {num_users} users with {messages_per_user} messages each...")
        
        for i in range(num_users):
            user_id = f"memory_test_user_{i}"
            session = await create_session(user_id=user_id)
            
            for j in range(messages_per_user):
                await track_message(
                    session_id=session.session_id,
                    prompt_length=50,
                    response_length=100,
                    response_time_ms=1000,
                    user_id=user_id
                )
        
        # Get final memory usage
        final_memory = process.memory_info().rss / 1024 / 1024  # MB
        memory_increase = final_memory - initial_memory
        
        print(f"✅ Memory usage: {initial_memory:.1f}MB → {final_memory:.1f}MB (+{memory_increase:.1f}MB)")
        
        # Memory increase should be reasonable
        total_records = num_users * (1 + messages_per_user)  # sessions + messages
        memory_per_record = memory_increase / total_records
        
        print(f"✅ Memory per record: {memory_per_record:.3f}MB")
        
        # Performance assertion (should be less than 1MB per record)
        assert memory_per_record < 1.0, f"Memory usage too high: {memory_per_record:.3f}MB per record"


async def run_performance_tests():
    """Run all performance tests"""
    print("⚡ Running Performance Tests for User Authentication")
    print("=" * 60)
    
    # Test database index performance
    index_test = TestDatabaseIndexPerformance()
    await index_test.test_user_id_index_performance()
    await index_test.test_compound_index_performance()
    await index_test.test_sparse_index_performance()
    print("✅ Database index performance tests completed")
    
    # Test analytics function performance
    analytics_test = TestAnalyticsFunctionPerformance()
    await analytics_test.test_user_statistics_performance()
    await analytics_test.test_user_analytics_performance()
    await analytics_test.test_authenticated_vs_anonymous_performance()
    await analytics_test.test_basic_stats_with_filter_performance()
    await analytics_test.test_hourly_stats_with_filter_performance()
    print("✅ Analytics function performance tests completed")
    
    # Test concurrent performance
    concurrent_test = TestConcurrentUserPerformance()
    await concurrent_test.test_concurrent_user_creation()
    await concurrent_test.test_concurrent_user_queries()
    print("✅ Concurrent performance tests completed")
    
    # Test memory performance
    memory_test = TestMemoryPerformance()
    await memory_test.test_memory_usage_with_users()
    print("✅ Memory performance tests completed")
    
    print("\n🎉 ALL PERFORMANCE TESTS COMPLETED!")
    print("User authentication feature maintains good performance characteristics.")


if __name__ == "__main__":
    asyncio.run(run_performance_tests())