"""
Advanced memory optimization systems for Sprint 4 performance improvements
"""
from typing import Dict, Any, Optional, List, Set, TypeVar
import threading
from collections import defaultdict
from .performance import get_memory_manager

T = TypeVar('T')


class StringInternPool:
    """
    High-performance string interning system.
    Converts repeated string comparisons into fast integer comparisons.
    """
    
    def __init__(self, max_size: int = 50000):
        self.max_size = max_size
        self._string_to_id: Dict[str, int] = {}
        self._id_to_string: Dict[int, str] = {}
        self._next_id = 1
        self._lock = threading.RLock()
        self._access_count: Dict[int, int] = defaultdict(int)
        
        # Statistics
        self._stats = {
            "total_interns": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "memory_saved_bytes": 0
        }
    
    def intern(self, string: str) -> int:
        """
        Intern a string and return its unique integer ID.
        Subsequent calls with the same string return the same ID.
        """
        if not isinstance(string, str):
            string = str(string)
        
        with self._lock:
            # Check if already interned
            if string in self._string_to_id:
                string_id = self._string_to_id[string]
                self._access_count[string_id] += 1
                self._stats["cache_hits"] += 1
                return string_id
            
            # Check if we need to cleanup
            if len(self._string_to_id) >= self.max_size:
                self._cleanup_least_used()
            
            # Intern new string
            string_id = self._next_id
            self._next_id += 1
            
            self._string_to_id[string] = string_id
            self._id_to_string[string_id] = string
            self._access_count[string_id] = 1
            
            self._stats["total_interns"] += 1
            self._stats["cache_misses"] += 1
            self._stats["memory_saved_bytes"] += len(string) * 8  # Estimate
            
            return string_id
    
    def get_string(self, string_id: int) -> Optional[str]:
        """Get string by its interned ID."""
        with self._lock:
            return self._id_to_string.get(string_id)
    
    def get_id(self, string: str) -> Optional[int]:
        """Get ID for a string if it's interned."""
        with self._lock:
            return self._string_to_id.get(string)
    
    def compare_fast(self, id1: int, id2: int) -> bool:
        """Fast integer comparison instead of string comparison."""
        return id1 == id2
    
    def _cleanup_least_used(self):
        """Remove least frequently used strings."""
        if not self._access_count:
            return
        
        # Remove 25% of least used strings
        cleanup_count = max(1, len(self._access_count) // 4)
        
        # Sort by access count and remove least used
        sorted_items = sorted(self._access_count.items(), key=lambda x: x[1])
        
        for string_id, _ in sorted_items[:cleanup_count]:
            string = self._id_to_string.get(string_id)
            if string:
                del self._string_to_id[string]
                del self._id_to_string[string_id]
                del self._access_count[string_id]
    
    def get_stats(self) -> Dict[str, Any]:
        """Get interning statistics."""
        with self._lock:
            total_requests = self._stats["cache_hits"] + self._stats["cache_misses"]
            
            stats = self._stats.copy()
            stats.update({
                "pool_size": len(self._string_to_id),
                "max_size": self.max_size,
                "utilization": len(self._string_to_id) / self.max_size,
                "hit_rate": self._stats["cache_hits"] / max(total_requests, 1),
                "next_id": self._next_id - 1
            })
            
            return stats
    
    def clear(self):
        """Clear the intern pool."""
        with self._lock:
            self._string_to_id.clear()
            self._id_to_string.clear()
            self._access_count.clear()
            self._next_id = 1


class AdvancedMemoryPool:
    """
    Memory pool for fast allocation of frequently used objects.
    Reduces garbage collection pressure and improves cache locality.
    """
    
    def __init__(self, object_type: type, initial_size: int = 100, max_size: int = 1000):
        self.object_type = object_type
        self.initial_size = initial_size
        self.max_size = max_size
        
        self._available_objects: List[Any] = []
        self._allocated_objects: Set[int] = set()
        self._lock = threading.Lock()
        
        # Pre-allocate initial objects
        self._preallocate(initial_size)
        
        # Statistics
        self._stats = {
            "total_allocations": 0,
            "pool_hits": 0,
            "pool_misses": 0,
            "current_size": len(self._available_objects),
            "peak_size": len(self._available_objects)
        }
    
    def _preallocate(self, count: int):
        """Pre-allocate objects to the pool."""
        for _ in range(count):
            try:
                obj = self.object_type()
                self._available_objects.append(obj)
            except Exception:
                break  # Can't create more objects
    
    def acquire(self) -> Any:
        """Acquire an object from the pool."""
        with self._lock:
            if self._available_objects:
                obj = self._available_objects.pop()
                self._allocated_objects.add(id(obj))
                self._stats["pool_hits"] += 1
                self._stats["total_allocations"] += 1
                return obj
            else:
                # Pool is empty, create new object
                try:
                    obj = self.object_type()
                    self._allocated_objects.add(id(obj))
                    self._stats["pool_misses"] += 1
                    self._stats["total_allocations"] += 1
                    return obj
                except Exception:
                    raise MemoryError(f"Cannot allocate {self.object_type}")
    
    def release(self, obj: Any):
        """Return an object to the pool."""
        if obj is None:
            return
        
        obj_id = id(obj)
        
        with self._lock:
            if obj_id not in self._allocated_objects:
                return  # Object wasn't allocated by this pool
            
            self._allocated_objects.remove(obj_id)
            
            # Return to pool if not at max capacity
            if len(self._available_objects) < self.max_size:
                # Reset object state if possible
                if hasattr(obj, '_reset'):
                    try:
                        obj._reset()
                    except Exception:
                        pass  # Reset failed, don't return to pool
                        return
                
                self._available_objects.append(obj)
                self._stats["current_size"] = len(self._available_objects)
                self._stats["peak_size"] = max(self._stats["peak_size"], len(self._available_objects))
    
    def get_stats(self) -> Dict[str, Any]:
        """Get pool statistics."""
        with self._lock:
            total_requests = self._stats["pool_hits"] + self._stats["pool_misses"]
            
            stats = self._stats.copy()
            stats.update({
                "hit_rate": self._stats["pool_hits"] / max(total_requests, 1),
                "allocated_objects": len(self._allocated_objects),
                "available_objects": len(self._available_objects),
                "max_size": self.max_size
            })
            
            return stats
    
    def clear(self):
        """Clear the pool."""
        with self._lock:
            self._available_objects.clear()
            self._allocated_objects.clear()


class QueryOptimizationEngine:
    """
    Advanced query optimization using database-style techniques.
    """
    
    def __init__(self):
        self._query_stats: Dict[str, Dict[str, Any]] = {}
        self._optimization_cache: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.RLock()
    
    def analyze_query_pattern(self, selector: str, result_count: int, total_nodes: int, execution_time: float):
        """Analyze query pattern for optimization opportunities."""
        with self._lock:
            if selector not in self._query_stats:
                self._query_stats[selector] = {
                    "execution_count": 0,
                    "total_time": 0.0,
                    "total_results": 0,
                    "avg_selectivity": 0.0,
                    "avg_execution_time": 0.0
                }
            
            stats = self._query_stats[selector]
            stats["execution_count"] += 1
            stats["total_time"] += execution_time
            stats["total_results"] += result_count
            
            # Update averages
            stats["avg_execution_time"] = stats["total_time"] / stats["execution_count"]
            if total_nodes > 0:
                selectivity = result_count / total_nodes
                stats["avg_selectivity"] = ((stats["avg_selectivity"] * (stats["execution_count"] - 1)) + selectivity) / stats["execution_count"]
    
    def get_optimization_hints(self, selector: str) -> Dict[str, Any]:
        """Get optimization hints for a selector."""
        with self._lock:
            if selector in self._optimization_cache:
                return self._optimization_cache[selector]
            
            hints = self._generate_optimization_hints(selector)
            self._optimization_cache[selector] = hints
            return hints
    
    def _generate_optimization_hints(self, selector: str) -> Dict[str, Any]:
        """Generate optimization hints based on query analysis."""
        hints = {
            "use_index": False,
            "reverse_search": False,
            "cache_results": False,
            "optimization_level": "basic"
        }
        
        # Check if we have statistics for this selector
        if selector in self._query_stats:
            stats = self._query_stats[selector]
            
            # High-frequency queries should be cached
            if stats["execution_count"] > 10:
                hints["cache_results"] = True
            
            # Low selectivity queries might benefit from indexing
            if stats["avg_selectivity"] < 0.1:
                hints["use_index"] = True
                hints["optimization_level"] = "advanced"
            
            # Slow queries need optimization
            if stats["avg_execution_time"] > 0.01:  # 10ms
                hints["optimization_level"] = "aggressive"
        
        # Analyze selector structure
        if ' ' in selector and not any(op in selector for op in ['>', '+', '~']):
            # Descendant selector - might benefit from reverse search
            parts = selector.split()
            if len(parts) > 2:
                hints["reverse_search"] = True
        
        return hints
    
    def get_query_stats(self) -> Dict[str, Any]:
        """Get comprehensive query statistics."""
        with self._lock:
            return {
                "total_unique_queries": len(self._query_stats),
                "cached_optimizations": len(self._optimization_cache),
                "query_details": self._query_stats.copy()
            }


class IndexManager:
    """
    Advanced indexing system for fast element lookups.
    """
    
    def __init__(self):
        self._tag_index: Dict[str, Set[int]] = defaultdict(set)
        self._class_index: Dict[str, Set[int]] = defaultdict(set)
        self._id_index: Dict[str, int] = {}
        self._attribute_index: Dict[str, Dict[str, Set[int]]] = defaultdict(lambda: defaultdict(set))
        
        self._lock = threading.RLock()
        self._dirty = True
        self._index_stats = {
            "total_indexed_nodes": 0,
            "index_builds": 0,
            "index_hits": 0,
            "index_misses": 0
        }
    
    def build_index(self, root_tag):
        """Build comprehensive indices for fast lookups."""
        with self._lock:
            # Clear existing indices
            self._tag_index.clear()
            self._class_index.clear()
            self._id_index.clear()
            self._attribute_index.clear()
            
            node_count = 0
            
            # Index root tag
            if hasattr(root_tag, 'name'):
                self._index_single_tag(root_tag, node_count)
                node_count += 1
            
            # Index all descendants
            try:
                for descendant in root_tag.descendants:
                    self._index_single_tag(descendant, node_count)
                    node_count += 1
            except (AttributeError, TypeError):
                pass  # Fallback if descendants fails
            
            self._index_stats["total_indexed_nodes"] = node_count
            self._index_stats["index_builds"] += 1
            self._dirty = False
    
    def _index_single_tag(self, tag, node_id: int):
        """Index a single tag."""
        # Index by tag name
        if hasattr(tag, 'name') and tag.name:
            self._tag_index[tag.name].add(node_id)
        
        # Index by attributes
        try:
            attrs = tag.attrs if hasattr(tag, 'attrs') else {}
            
            # Index ID
            if 'id' in attrs and attrs['id']:
                self._id_index[attrs['id']] = node_id
            
            # Index classes
            if 'class' in attrs and attrs['class']:
                classes = attrs['class'].split()
                for class_name in classes:
                    if class_name:
                        self._class_index[class_name].add(node_id)
            
            # Index other attributes
            for attr_name, attr_value in attrs.items():
                if attr_name not in ['id', 'class'] and attr_value:
                    self._attribute_index[attr_name][attr_value].add(node_id)
        
        except (AttributeError, KeyError, TypeError):
            pass  # Skip if attribute access fails
    
    def lookup_by_tag(self, tag_name: str) -> Set[int]:
        """Fast lookup by tag name."""
        with self._lock:
            if self._dirty:
                self._index_stats["index_misses"] += 1
                return set()
            
            result = self._tag_index.get(tag_name, set())
            if result:
                self._index_stats["index_hits"] += 1
            else:
                self._index_stats["index_misses"] += 1
            
            return result.copy()
    
    def lookup_by_class(self, class_name: str) -> Set[int]:
        """Fast lookup by class name."""
        with self._lock:
            if self._dirty:
                self._index_stats["index_misses"] += 1
                return set()
            
            result = self._class_index.get(class_name, set())
            if result:
                self._index_stats["index_hits"] += 1
            else:
                self._index_stats["index_misses"] += 1
            
            return result.copy()
    
    def lookup_by_id(self, id_name: str) -> Optional[int]:
        """Fast lookup by ID."""
        with self._lock:
            if self._dirty:
                self._index_stats["index_misses"] += 1
                return None
            
            result = self._id_index.get(id_name)
            if result is not None:
                self._index_stats["index_hits"] += 1
            else:
                self._index_stats["index_misses"] += 1
            
            return result
    
    def lookup_by_attribute(self, attr_name: str, attr_value: str) -> Set[int]:
        """Fast lookup by attribute."""
        with self._lock:
            if self._dirty:
                self._index_stats["index_misses"] += 1
                return set()
            
            result = self._attribute_index.get(attr_name, {}).get(attr_value, set())
            if result:
                self._index_stats["index_hits"] += 1
            else:
                self._index_stats["index_misses"] += 1
            
            return result.copy()
    
    def invalidate(self):
        """Mark indices as dirty (need rebuild)."""
        with self._lock:
            self._dirty = True
    
    def get_stats(self) -> Dict[str, Any]:
        """Get indexing statistics."""
        with self._lock:
            total_lookups = self._index_stats["index_hits"] + self._index_stats["index_misses"]
            
            stats = self._index_stats.copy()
            stats.update({
                "tag_index_size": len(self._tag_index),
                "class_index_size": len(self._class_index),
                "id_index_size": len(self._id_index),
                "attribute_indices": len(self._attribute_index),
                "hit_rate": self._index_stats["index_hits"] / max(total_lookups, 1),
                "is_dirty": self._dirty
            })
            
            return stats


# Global instances
_global_string_pool: Optional[StringInternPool] = None
_global_optimization_engine: Optional[QueryOptimizationEngine] = None
_global_index_manager: Optional[IndexManager] = None


def get_string_pool() -> StringInternPool:
    """Get global string interning pool."""
    global _global_string_pool
    if _global_string_pool is None:
        _global_string_pool = StringInternPool()
    return _global_string_pool


def get_optimization_engine() -> QueryOptimizationEngine:
    """Get global query optimization engine."""
    global _global_optimization_engine
    if _global_optimization_engine is None:
        _global_optimization_engine = QueryOptimizationEngine()
    return _global_optimization_engine


def get_index_manager() -> IndexManager:
    """Get global index manager."""
    global _global_index_manager
    if _global_index_manager is None:
        _global_index_manager = IndexManager()
    return _global_index_manager


def intern_string(string: str) -> int:
    """Convenience function to intern a string."""
    return get_string_pool().intern(string)


def compare_interned(id1: int, id2: int) -> bool:
    """Convenience function for fast string comparison."""
    return get_string_pool().compare_fast(id1, id2)


def optimize_memory_globally():
    """Perform global memory optimization."""
    # Get all global instances and optimize
    memory_mgr = get_memory_manager()
    memory_mgr.cleanup(force=True)
    
    # String pool has internal cleanup
    get_string_pool()  # Trigger any internal cleanup
    
    # Force garbage collection
    import gc
    gc.collect()


def get_memory_optimization_stats() -> Dict[str, Any]:
    """Get comprehensive memory optimization statistics."""
    return {
        "string_pool": get_string_pool().get_stats(),
        "optimization_engine": get_optimization_engine().get_query_stats(),
        "index_manager": get_index_manager().get_stats(),
        "memory_manager": get_memory_manager().get_memory_stats()
    }