"""
Performance optimization utilities for Toxo.

This module provides performance optimizations including:
- Memory pooling for embeddings
- Response caching mechanisms
- Vector similarity search optimization
- Memory usage monitoring and cleanup
"""

import gc
import time
import asyncio
import threading
import psutil
import numpy as np
from typing import Dict, Any, Optional, List, Tuple, Callable, Union
from collections import OrderedDict, defaultdict
from dataclasses import dataclass, field
from functools import wraps, lru_cache
import logging
import weakref
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager

from .exceptions import MemoryError as ToxoMemoryError, PerformanceWarning

logger = logging.getLogger(__name__)


@dataclass
class PerformanceMetrics:
    """Performance metrics tracking."""
    operation_count: int = 0
    total_time: float = 0.0
    average_time: float = 0.0
    min_time: float = float('inf')
    max_time: float = 0.0
    memory_usage_mb: float = 0.0
    cache_hits: int = 0
    cache_misses: int = 0
    
    def update(self, duration: float, memory_mb: float = 0.0):
        """Update metrics with new operation."""
        self.operation_count += 1
        self.total_time += duration
        self.average_time = self.total_time / self.operation_count
        self.min_time = min(self.min_time, duration)
        self.max_time = max(self.max_time, duration)
        if memory_mb > 0:
            self.memory_usage_mb = memory_mb


class MemoryPool:
    """Memory pool for reusing embedding vectors and large objects."""
    
    def __init__(self, pool_size: int = 1000, cleanup_threshold: float = 0.8):
        """
        Initialize memory pool.
        
        Args:
            pool_size: Maximum number of objects to pool
            cleanup_threshold: Memory usage threshold for cleanup (0.0-1.0)
        """
        self.pool_size = pool_size
        self.cleanup_threshold = cleanup_threshold
        self._pools: Dict[str, List[Any]] = defaultdict(list)
        self._usage_count: Dict[str, int] = defaultdict(int)
        self._lock = threading.RLock()
        self._total_allocated = 0
        
    def get_embedding_vector(self, dimension: int, dtype: type = np.float32) -> np.ndarray:
        """Get a reusable embedding vector from pool."""
        pool_key = f"embedding_{dimension}_{dtype.__name__}"
        
        with self._lock:
            pool = self._pools[pool_key]
            if pool:
                vector = pool.pop()
                vector.fill(0)  # Reset to zeros
                self._usage_count[pool_key] += 1
                return vector
            else:
                # Create new vector if pool is empty
                vector = np.zeros(dimension, dtype=dtype)
                self._total_allocated += vector.nbytes
                self._usage_count[pool_key] += 1
                
                # Check if we need cleanup
                if self._should_cleanup():
                    self._cleanup_pools()
                
                return vector
    
    def return_embedding_vector(self, vector: np.ndarray):
        """Return a vector to the pool for reuse."""
        if vector is None:
            return
        
        pool_key = f"embedding_{vector.shape[0]}_{vector.dtype.name}"
        
        with self._lock:
            pool = self._pools[pool_key]
            if len(pool) < self.pool_size:
                pool.append(vector)
    
    def get_matrix(self, shape: Tuple[int, ...], dtype: type = np.float32) -> np.ndarray:
        """Get a reusable matrix from pool."""
        pool_key = f"matrix_{'x'.join(map(str, shape))}_{dtype.__name__}"
        
        with self._lock:
            pool = self._pools[pool_key]
            if pool:
                matrix = pool.pop()
                matrix.fill(0)
                self._usage_count[pool_key] += 1
                return matrix
            else:
                matrix = np.zeros(shape, dtype=dtype)
                self._total_allocated += matrix.nbytes
                self._usage_count[pool_key] += 1
                return matrix
    
    def return_matrix(self, matrix: np.ndarray):
        """Return a matrix to the pool for reuse."""
        if matrix is None:
            return
        
        pool_key = f"matrix_{'x'.join(map(str, matrix.shape))}_{matrix.dtype.name}"
        
        with self._lock:
            pool = self._pools[pool_key]
            if len(pool) < self.pool_size:
                pool.append(matrix)
    
    def _should_cleanup(self) -> bool:
        """Check if cleanup is needed based on memory usage."""
        current_memory = get_memory_usage_mb()
        system_memory = psutil.virtual_memory().total / (1024 * 1024)
        usage_ratio = current_memory / system_memory
        return usage_ratio > self.cleanup_threshold
    
    def _cleanup_pools(self):
        """Clean up least used pools to free memory."""
        logger.info("Starting memory pool cleanup")
        
        # Sort pools by usage count (ascending)
        sorted_pools = sorted(self._usage_count.items(), key=lambda x: x[1])
        
        cleaned_count = 0
        for pool_key, _ in sorted_pools[:len(sorted_pools) // 2]:
            pool = self._pools[pool_key]
            cleaned_count += len(pool)
            pool.clear()
        
        # Force garbage collection
        gc.collect()
        
        logger.info(f"Cleaned {cleaned_count} objects from memory pools")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get memory pool statistics."""
        with self._lock:
            return {
                "total_pools": len(self._pools),
                "total_objects": sum(len(pool) for pool in self._pools.values()),
                "pool_usage": dict(self._usage_count),
                "estimated_memory_mb": self._total_allocated / (1024 * 1024)
            }


class LRUCache:
    """Enhanced LRU cache with size limits and TTL."""
    
    def __init__(self, maxsize: int = 1024, ttl: Optional[float] = None):
        """
        Initialize LRU cache.
        
        Args:
            maxsize: Maximum number of items to cache
            ttl: Time-to-live for cache entries in seconds
        """
        self.maxsize = maxsize
        self.ttl = ttl
        self._cache: OrderedDict = OrderedDict()
        self._timestamps: Dict[str, float] = {} if ttl else None
        self._lock = threading.RLock()
        self._hits = 0
        self._misses = 0
    
    def get(self, key: str) -> Any:
        """Get item from cache."""
        with self._lock:
            # Check if key exists and is not expired
            if key in self._cache:
                if self._is_expired(key):
                    self._remove(key)
                    self._misses += 1
                    return None
                
                # Move to end (most recently used)
                value = self._cache.pop(key)
                self._cache[key] = value
                self._hits += 1
                return value
            
            self._misses += 1
            return None
    
    def put(self, key: str, value: Any):
        """Put item in cache."""
        with self._lock:
            # Remove if already exists
            if key in self._cache:
                self._remove(key)
            
            # Check size limit
            if len(self._cache) >= self.maxsize:
                # Remove least recently used
                oldest_key = next(iter(self._cache))
                self._remove(oldest_key)
            
            # Add new item
            self._cache[key] = value
            if self._timestamps is not None:
                self._timestamps[key] = time.time()
    
    def _remove(self, key: str):
        """Remove item from cache."""
        if key in self._cache:
            del self._cache[key]
        if self._timestamps and key in self._timestamps:
            del self._timestamps[key]
    
    def _is_expired(self, key: str) -> bool:
        """Check if cache entry is expired."""
        if self.ttl is None or self._timestamps is None:
            return False
        
        timestamp = self._timestamps.get(key)
        if timestamp is None:
            return True
        
        return time.time() - timestamp > self.ttl
    
    def clear(self):
        """Clear all cache entries."""
        with self._lock:
            self._cache.clear()
            if self._timestamps:
                self._timestamps.clear()
    
    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics."""
        with self._lock:
            total_requests = self._hits + self._misses
            hit_rate = self._hits / total_requests if total_requests > 0 else 0.0
            
            return {
                "size": len(self._cache),
                "maxsize": self.maxsize,
                "hits": self._hits,
                "misses": self._misses,
                "hit_rate": hit_rate,
                "ttl": self.ttl
            }


class VectorSearchOptimizer:
    """Optimized vector similarity search with indexing."""
    
    def __init__(self, dimension: int, index_type: str = "flat"):
        """
        Initialize vector search optimizer.
        
        Args:
            dimension: Vector dimension
            index_type: Index type ("flat", "ivf", "hnsw")
        """
        self.dimension = dimension
        self.index_type = index_type
        self._vectors: List[np.ndarray] = []
        self._ids: List[str] = []
        self._id_to_index: Dict[str, int] = {}
        self._index_built = False
        self._lock = threading.RLock()
        
        # Initialize based on index type
        if index_type == "flat":
            self._index = None  # Use numpy operations
        else:
            # For advanced indexing, would use libraries like faiss
            logger.warning(f"Index type {index_type} not fully implemented, using flat")
            self._index = None
    
    def add_vector(self, vector_id: str, vector: np.ndarray):
        """Add vector to index."""
        if vector.shape[0] != self.dimension:
            raise ValueError(f"Vector dimension {vector.shape[0]} doesn't match expected {self.dimension}")
        
        with self._lock:
            if vector_id in self._id_to_index:
                # Update existing vector
                index = self._id_to_index[vector_id]
                self._vectors[index] = vector.copy()
            else:
                # Add new vector
                self._vectors.append(vector.copy())
                index = len(self._vectors) - 1
                self._ids.append(vector_id)
                self._id_to_index[vector_id] = index
            
            self._index_built = False  # Rebuild index on next search
    
    def search(self, query_vector: np.ndarray, k: int = 10, threshold: float = 0.0) -> List[Tuple[str, float]]:
        """
        Search for similar vectors.
        
        Args:
            query_vector: Query vector
            k: Number of results to return
            threshold: Similarity threshold
            
        Returns:
            List of (vector_id, similarity_score) tuples
        """
        if query_vector.shape[0] != self.dimension:
            raise ValueError(f"Query vector dimension {query_vector.shape[0]} doesn't match expected {self.dimension}")
        
        with self._lock:
            if not self._vectors:
                return []
            
            # Build index if needed
            if not self._index_built:
                self._build_index()
            
            # Compute similarities
            if self._index is None:
                # Flat search using numpy
                vectors_matrix = np.vstack(self._vectors)
                similarities = self._compute_cosine_similarities(query_vector, vectors_matrix)
            else:
                # Use advanced index
                similarities = self._index_search(query_vector)
            
            # Filter by threshold and get top-k
            results = []
            for i, similarity in enumerate(similarities):
                if similarity >= threshold:
                    results.append((self._ids[i], similarity))
            
            # Sort by similarity (descending) and take top-k
            results.sort(key=lambda x: x[1], reverse=True)
            return results[:k]
    
    def _build_index(self):
        """Build search index."""
        # For flat index, no building needed
        if self.index_type == "flat":
            self._index_built = True
            return
        
        # For advanced indexes, build here
        logger.info(f"Building {self.index_type} index for {len(self._vectors)} vectors")
        self._index_built = True
    
    def _compute_cosine_similarities(self, query: np.ndarray, vectors: np.ndarray) -> np.ndarray:
        """Compute cosine similarities efficiently."""
        # Normalize vectors
        query_norm = query / (np.linalg.norm(query) + 1e-8)
        vectors_norm = vectors / (np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-8)
        
        # Compute dot product (cosine similarity for normalized vectors)
        similarities = np.dot(vectors_norm, query_norm)
        return similarities
    
    def _index_search(self, query_vector: np.ndarray) -> np.ndarray:
        """Search using built index."""
        # Placeholder for advanced index search
        vectors_matrix = np.vstack(self._vectors)
        return self._compute_cosine_similarities(query_vector, vectors_matrix)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get search optimizer statistics."""
        with self._lock:
            return {
                "vector_count": len(self._vectors),
                "dimension": self.dimension,
                "index_type": self.index_type,
                "index_built": self._index_built,
                "memory_mb": sum(v.nbytes for v in self._vectors) / (1024 * 1024)
            }


class MemoryMonitor:
    """Monitor and manage memory usage."""
    
    def __init__(self, warning_threshold: float = 0.8, critical_threshold: float = 0.9):
        """
        Initialize memory monitor.
        
        Args:
            warning_threshold: Memory usage ratio to trigger warnings
            critical_threshold: Memory usage ratio to trigger cleanup
        """
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self._callbacks: List[Callable] = []
        self._monitoring = False
        self._monitor_thread = None
        
    def add_cleanup_callback(self, callback: Callable):
        """Add callback to be called when memory cleanup is needed."""
        self._callbacks.append(callback)
    
    def start_monitoring(self, interval: float = 5.0):
        """Start memory monitoring in background thread."""
        if self._monitoring:
            return
        
        self._monitoring = True
        self._monitor_thread = threading.Thread(
            target=self._monitor_loop,
            args=(interval,),
            daemon=True
        )
        self._monitor_thread.start()
        logger.info("Memory monitoring started")
    
    def stop_monitoring(self):
        """Stop memory monitoring."""
        self._monitoring = False
        if self._monitor_thread:
            self._monitor_thread.join(timeout=1.0)
        logger.info("Memory monitoring stopped")
    
    def _monitor_loop(self, interval: float):
        """Main monitoring loop."""
        while self._monitoring:
            try:
                memory_info = self.get_memory_info()
                usage_ratio = memory_info['usage_ratio']
                
                if usage_ratio >= self.critical_threshold:
                    logger.critical(f"Critical memory usage: {usage_ratio:.1%}")
                    self._trigger_cleanup()
                elif usage_ratio >= self.warning_threshold:
                    logger.warning(f"High memory usage: {usage_ratio:.1%}")
                
                time.sleep(interval)
            except Exception as e:
                logger.error(f"Error in memory monitoring: {e}")
                time.sleep(interval)
    
    def _trigger_cleanup(self):
        """Trigger cleanup callbacks."""
        for callback in self._callbacks:
            try:
                callback()
            except Exception as e:
                logger.error(f"Error in cleanup callback: {e}")
        
        # Force garbage collection
        gc.collect()
    
    def get_memory_info(self) -> Dict[str, Any]:
        """Get current memory information."""
        process = psutil.Process()
        memory_info = process.memory_info()
        virtual_memory = psutil.virtual_memory()
        
        return {
            "process_memory_mb": memory_info.rss / (1024 * 1024),
            "virtual_memory_mb": memory_info.vms / (1024 * 1024),
            "system_total_mb": virtual_memory.total / (1024 * 1024),
            "system_available_mb": virtual_memory.available / (1024 * 1024),
            "system_used_mb": virtual_memory.used / (1024 * 1024),
            "usage_ratio": virtual_memory.percent / 100.0
        }


class PerformanceTracker:
    """Advanced performance tracking and monitoring."""
    
    def __init__(self, max_metrics: int = 1000):
        self.max_metrics = max_metrics
        self._metrics: Dict[str, PerformanceMetrics] = defaultdict(PerformanceMetrics)
        self._start_times: Dict[str, float] = {}
        self._lock = threading.Lock()
    
    @contextmanager
    def track(self, operation_name: str):
        """Context manager for tracking operation performance."""
        start_time = time.time()
        start_memory = self._get_memory_usage()
        
        try:
            yield
        finally:
            end_time = time.time()
            end_memory = self._get_memory_usage()
            
            duration = end_time - start_time
            memory_delta = max(0, end_memory - start_memory)
            
            with self._lock:
                self._metrics[operation_name].update(duration, memory_delta)
    
    def start_operation(self, operation_name: str):
        """Start tracking an operation."""
        with self._lock:
            self._start_times[operation_name] = time.time()
    
    def end_operation(self, operation_name: str):
        """End tracking an operation."""
        end_time = time.time()
        
        with self._lock:
            start_time = self._start_times.pop(operation_name, end_time)
            duration = end_time - start_time
            self._metrics[operation_name].update(duration)
    
    def record_cache_hit(self, operation_name: str):
        """Record a cache hit."""
        with self._lock:
            self._metrics[operation_name].cache_hits += 1
    
    def record_cache_miss(self, operation_name: str):
        """Record a cache miss."""
        with self._lock:
            self._metrics[operation_name].cache_misses += 1
    
    def get_summary(self) -> Dict[str, Any]:
        """Get performance summary."""
        with self._lock:
            summary = {}
            for op_name, metrics in self._metrics.items():
                summary[op_name] = {
                    "operation_count": metrics.operation_count,
                    "total_time": metrics.total_time,
                    "average_time": metrics.average_time,
                    "min_time": metrics.min_time if metrics.min_time != float('inf') else 0,
                    "max_time": metrics.max_time,
                    "memory_usage_mb": metrics.memory_usage_mb,
                    "cache_hit_rate": (
                        metrics.cache_hits / (metrics.cache_hits + metrics.cache_misses)
                        if (metrics.cache_hits + metrics.cache_misses) > 0 else 0
                    )
                }
            return summary
    
    def clear(self):
        """Clear all metrics."""
        with self._lock:
            self._metrics.clear()
            self._start_times.clear()
    
    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB."""
        try:
            process = psutil.Process()
            return process.memory_info().rss / (1024 * 1024)
        except:
            return 0.0


# Global instances
_memory_pool = MemoryPool()
_response_cache = LRUCache(maxsize=1000, ttl=3600)  # 1 hour TTL
_memory_monitor = MemoryMonitor()
_performance_metrics: Dict[str, PerformanceMetrics] = defaultdict(PerformanceMetrics)
_performance_tracker = PerformanceTracker()


def get_memory_pool() -> MemoryPool:
    """Get global memory pool instance."""
    return _memory_pool


def get_response_cache() -> LRUCache:
    """Get global response cache instance."""
    return _response_cache


def get_memory_monitor() -> MemoryMonitor:
    """Get global memory monitor instance."""
    return _memory_monitor


def performance_monitor(operation_name: str = None):
    """Decorator to monitor performance of functions."""
    def decorator(func):
        op_name = operation_name or f"{func.__module__}.{func.__name__}"
        
        @wraps(func)
        def wrapper(*args, **kwargs):
            start_time = time.time()
            start_memory = get_memory_usage_mb()
            
            try:
                result = func(*args, **kwargs)
                return result
            finally:
                end_time = time.time()
                end_memory = get_memory_usage_mb()
                
                duration = end_time - start_time
                memory_delta = end_memory - start_memory
                
                _performance_metrics[op_name].update(duration, max(0, memory_delta))
        
        return wrapper
    return decorator


def cache_response(cache_key_func: Callable = None, ttl: Optional[float] = None):
    """Decorator to cache function responses."""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Generate cache key
            if cache_key_func:
                cache_key = cache_key_func(*args, **kwargs)
            else:
                cache_key = f"{func.__name__}:{hash(str(args) + str(sorted(kwargs.items())))}"
            
            # Try to get from cache
            cached_result = _response_cache.get(cache_key)
            if cached_result is not None:
                return cached_result
            
            # Compute result and cache it
            result = func(*args, **kwargs)
            _response_cache.put(cache_key, result)
            
            return result
        
        return wrapper
    return decorator


def get_memory_usage_mb() -> float:
    """Get current process memory usage in MB."""
    process = psutil.Process()
    return process.memory_info().rss / (1024 * 1024)


def get_performance_stats() -> Dict[str, Dict[str, Any]]:
    """Get performance statistics for all monitored operations."""
    stats = {}
    for op_name, metrics in _performance_metrics.items():
        stats[op_name] = {
            "operation_count": metrics.operation_count,
            "total_time": metrics.total_time,
            "average_time": metrics.average_time,
            "min_time": metrics.min_time if metrics.min_time != float('inf') else 0,
            "max_time": metrics.max_time,
            "memory_usage_mb": metrics.memory_usage_mb,
            "cache_hit_rate": metrics.cache_hits / (metrics.cache_hits + metrics.cache_misses) 
                             if (metrics.cache_hits + metrics.cache_misses) > 0 else 0
        }
    
    return stats


def cleanup_performance_data():
    """Clean up performance monitoring data."""
    _performance_metrics.clear()
    _response_cache.clear()
    _memory_pool._cleanup_pools()
    gc.collect()


# Initialize memory monitoring
def start_performance_monitoring():
    """Start performance monitoring systems."""
    _memory_monitor.add_cleanup_callback(cleanup_performance_data)
    _memory_monitor.start_monitoring()


def stop_performance_monitoring():
    """Stop performance monitoring systems."""
    _memory_monitor.stop_monitoring() 