#!/usr/bin/env python3
"""
Benchmark script for Hybrid Search performance

This script tests the performance of different search strategies
to ensure they meet the targets specified in Phase 2.
"""

import time
import asyncio
import json
import numpy as np
from pathlib import Path
import sys

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from core.search.hybrid_search_service import HybridSearchService
from core.search.query_strategy_router import QueryStrategyRouter
from core.intelligence.conversation_intelligence import ConversationIntelligence
from core.storage.chroma_client import get_client as get_chroma_client


class SearchBenchmark:
    """Benchmark harness for search performance testing."""
    
    def __init__(self):
        """Initialize benchmark components."""
        self.hybrid_search = HybridSearchService()
        self.router = QueryStrategyRouter()
        self.results = {
            'simple_queries': [],
            'complex_queries': [],
            'hybrid_queries': [],
            'metadata_queries': []
        }
    
    async def run_benchmarks(self):
        """Run all benchmark tests."""
        print("🚀 Starting Hybrid Search Benchmarks")
        print("=" * 60)
        
        # Prepare test data
        await self._prepare_test_data()
        
        # Run query benchmarks
        await self._benchmark_simple_queries()
        await self._benchmark_complex_queries()
        await self._benchmark_hybrid_queries()
        await self._benchmark_metadata_queries()
        
        # Generate report
        self._generate_report()
    
    async def _prepare_test_data(self):
        """Prepare test data in ChromaDB."""
        print("\n📊 Preparing test data...")
        
        # Initialize ChromaDB
        chroma_client = get_chroma_client()
        chroma_client.initialize_collections()
        
        # Add test conversations
        conv_intel = ConversationIntelligence()
        
        test_conversations = [
            {
                "id": f"test_conv_{i}",
                "timestamp": f"2024-01-{i:02d}T10:00:00Z",
                "messages": [
                    {"role": "user", "content": f"Test message about {topic}"},
                    {"role": "assistant", "content": f"Response about {topic} implementation"}
                ],
                "project": f"project_{i % 3}",
                "topics": [topic]
            }
            for i, topic in enumerate([
                "neural networks", "memory systems", "search optimization",
                "pattern recognition", "consciousness modeling", "FAISS integration",
                "ChromaDB setup", "hybrid architecture", "performance tuning",
                "quantum entanglement"
            ], 1)
        ]
        
        for conv in test_conversations:
            conv_intel.analyze_conversation(conv)
        
        print(f"✅ Added {len(test_conversations)} test conversations")
    
    async def _benchmark_simple_queries(self):
        """Benchmark simple keyword queries (target: <100ms)."""
        print("\n🔍 Benchmarking simple queries...")
        
        simple_queries = [
            "neural",
            "memory",
            "FAISS",
            "search optimization",
            "pattern"
        ]
        
        for query in simple_queries:
            start_time = time.time()
            
            # Route query
            routing = self.router.route_query(query)
            
            # Execute search
            results = await self.hybrid_search.intelligent_search(query, top_k=5)
            
            elapsed = (time.time() - start_time) * 1000  # Convert to ms
            
            self.results['simple_queries'].append({
                'query': query,
                'time_ms': elapsed,
                'strategy': routing['strategy'],
                'result_count': len(results)
            })
            
            print(f"  ✓ '{query}': {elapsed:.1f}ms ({routing['strategy']})")
    
    async def _benchmark_complex_queries(self):
        """Benchmark complex semantic queries (target: <2000ms)."""
        print("\n🧠 Benchmarking complex queries...")
        
        complex_queries = [
            "how does neural network integration work",
            "what is the pattern recognition system",
            "explain consciousness modeling approach",
            "why use hybrid search architecture",
            "describe memory system optimization"
        ]
        
        for query in complex_queries:
            start_time = time.time()
            
            # Route query
            routing = self.router.route_query(query)
            
            # Execute search
            results = await self.hybrid_search.intelligent_search(query, top_k=10)
            
            elapsed = (time.time() - start_time) * 1000
            
            self.results['complex_queries'].append({
                'query': query,
                'time_ms': elapsed,
                'strategy': routing['strategy'],
                'result_count': len(results)
            })
            
            print(f"  ✓ '{query[:40]}...': {elapsed:.1f}ms ({routing['strategy']})")
    
    async def _benchmark_hybrid_queries(self):
        """Benchmark hybrid queries (target: <1000ms)."""
        print("\n🔄 Benchmarking hybrid queries...")
        
        hybrid_queries = [
            "find patterns similar to neural networks",
            "search memory optimization but not FAISS",
            "show recent consciousness modeling updates",
            "get performance tuning like search optimization"
        ]
        
        for query in hybrid_queries:
            start_time = time.time()
            
            # Route query
            routing = self.router.route_query(query)
            
            # Execute search
            results = await self.hybrid_search.intelligent_search(query, top_k=10)
            
            elapsed = (time.time() - start_time) * 1000
            
            self.results['hybrid_queries'].append({
                'query': query,
                'time_ms': elapsed,
                'strategy': routing['strategy'],
                'result_count': len(results)
            })
            
            print(f"  ✓ '{query[:40]}...': {elapsed:.1f}ms ({routing['strategy']})")
    
    async def _benchmark_metadata_queries(self):
        """Benchmark metadata-filtered queries."""
        print("\n🏷️ Benchmarking metadata queries...")
        
        metadata_queries = [
            "project:project_1 neural",
            "complexity:high optimization",
            "sentiment:positive implementation",
            "after:2024-01-05 search"
        ]
        
        for query in metadata_queries:
            start_time = time.time()
            
            # Route query
            routing = self.router.route_query(query)
            
            # Execute search with parsed context
            context = self._parse_metadata_query(query)
            results = await self.hybrid_search.intelligent_search(
                context['query'], 
                context=context.get('filters'),
                top_k=10
            )
            
            elapsed = (time.time() - start_time) * 1000
            
            self.results['metadata_queries'].append({
                'query': query,
                'time_ms': elapsed,
                'strategy': routing['strategy'],
                'result_count': len(results)
            })
            
            print(f"  ✓ '{query}': {elapsed:.1f}ms ({routing['strategy']})")
    
    def _parse_metadata_query(self, query: str) -> dict:
        """Parse metadata filters from query."""
        import re
        
        # Extract filters
        filters = {}
        clean_query = query
        
        # Project filter
        project_match = re.search(r'project:(\S+)', query)
        if project_match:
            filters['project'] = project_match.group(1)
            clean_query = clean_query.replace(project_match.group(0), '')
        
        # Complexity filter
        complexity_match = re.search(r'complexity:(\S+)', query)
        if complexity_match:
            filters['min_complexity'] = complexity_match.group(1)
            clean_query = clean_query.replace(complexity_match.group(0), '')
        
        # Sentiment filter
        sentiment_match = re.search(r'sentiment:(\S+)', query)
        if sentiment_match:
            filters['sentiment'] = sentiment_match.group(1)
            clean_query = clean_query.replace(sentiment_match.group(0), '')
        
        # Date filter
        after_match = re.search(r'after:(\S+)', query)
        if after_match:
            filters['start_date'] = after_match.group(1)
            clean_query = clean_query.replace(after_match.group(0), '')
        
        return {
            'query': clean_query.strip(),
            'filters': filters if filters else None
        }
    
    def _generate_report(self):
        """Generate performance report."""
        print("\n" + "=" * 60)
        print("📊 PERFORMANCE REPORT")
        print("=" * 60)
        
        # Performance targets
        targets = {
            'simple_queries': 100,    # <100ms
            'complex_queries': 2000,  # <2s
            'hybrid_queries': 1000,   # <1s
            'metadata_queries': 2000  # <2s
        }
        
        for query_type, results in self.results.items():
            if not results:
                continue
                
            times = [r['time_ms'] for r in results]
            avg_time = np.mean(times)
            max_time = np.max(times)
            min_time = np.min(times)
            target = targets.get(query_type, 1000)
            
            # Check if target met
            target_met = avg_time <= target
            status = "✅ PASS" if target_met else "❌ FAIL"
            
            print(f"\n{query_type.replace('_', ' ').title()}:")
            print(f"  Target:  <{target}ms")
            print(f"  Average: {avg_time:.1f}ms {status}")
            print(f"  Min:     {min_time:.1f}ms")
            print(f"  Max:     {max_time:.1f}ms")
            
            # Strategy distribution
            strategies = {}
            for r in results:
                strategy = r['strategy']
                strategies[strategy] = strategies.get(strategy, 0) + 1
            
            print(f"  Strategies: {dict(strategies)}")
        
        # Router statistics
        print("\n📈 Router Statistics:")
        router_stats = self.router.get_routing_stats()
        for strategy, stats in router_stats.items():
            if stats['count'] > 0:
                print(f"  {strategy}:")
                print(f"    Count: {stats['count']}")
                print(f"    Avg Confidence: {stats['avg_confidence']:.2f}")
                print(f"    Avg Complexity: {stats['avg_complexity']:.2f}")
        
        # Save detailed results
        results_file = Path(__file__).parent / "benchmark_results.json"
        with open(results_file, 'w') as f:
            json.dump({
                'timestamp': time.strftime('%Y-%m-%d %H:%M:%S'),
                'results': self.results,
                'router_stats': router_stats
            }, f, indent=2)
        
        print(f"\n💾 Detailed results saved to: {results_file}")


async def main():
    """Run benchmarks."""
    benchmark = SearchBenchmark()
    await benchmark.run_benchmarks()


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