"""
AbuLang Memory Manager - Memory optimization and management
"""

import gc
import sys
import psutil
import os
from typing import Dict, Any, List
from collections import defaultdict


class MemoryTracker:
    """Tracks memory usage"""
    
    def __init__(self):
        self.process = psutil.Process(os.getpid())
        self.snapshots: List[Dict] = []
    
    def get_memory_usage(self) -> Dict:
        """Get current memory usage"""
        mem_info = self.process.memory_info()
        return {
            "rss": mem_info.rss,  # Resident Set Size
            "vms": mem_info.vms,  # Virtual Memory Size
            "percent": self.process.memory_percent()
        }
    
    def take_snapshot(self, label: str = "") -> Dict:
        """Take memory snapshot"""
        snapshot = {
            "label": label,
            "memory": self.get_memory_usage(),
            "objects": len(gc.get_objects())
        }
        self.snapshots.append(snapshot)
        return snapshot
    
    def get_memory_delta(self, index1: int, index2: int) -> Dict:
        """Get memory difference between two snapshots"""
        if index1 >= len(self.snapshots) or index2 >= len(self.snapshots):
            return None
        
        snap1 = self.snapshots[index1]
        snap2 = self.snapshots[index2]
        
        return {
            "rss_delta": snap2["memory"]["rss"] - snap1["memory"]["rss"],
            "vms_delta": snap2["memory"]["vms"] - snap1["memory"]["vms"],
            "objects_delta": snap2["objects"] - snap1["objects"]
        }
    
    def get_report(self) -> str:
        """Get memory usage report"""
        current = self.get_memory_usage()
        report = "\n=== Memory Usage Report ===\n"
        report += f"RSS: {current['rss'] / 1024 / 1024:.2f} MB\n"
        report += f"VMS: {current['vms'] / 1024 / 1024:.2f} MB\n"
        report += f"Percent: {current['percent']:.2f}%\n"
        report += f"Objects: {len(gc.get_objects())}\n"
        return report


class MemoryOptimizer:
    """Optimizes memory usage"""
    
    def __init__(self):
        self.tracker = MemoryTracker()
        self.gc_enabled = True
    
    def enable_gc(self):
        """Enable garbage collection"""
        gc.enable()
        self.gc_enabled = True
    
    def disable_gc(self):
        """Disable garbage collection"""
        gc.disable()
        self.gc_enabled = False
    
    def collect(self) -> int:
        """Force garbage collection"""
        return gc.collect()
    
    def get_object_sizes(self) -> Dict[str, int]:
        """Get sizes of different object types"""
        sizes = defaultdict(int)
        for obj in gc.get_objects():
            obj_type = type(obj).__name__
            sizes[obj_type] += sys.getsizeof(obj)
        return dict(sizes)
    
    def find_memory_leaks(self) -> List[str]:
        """Find potential memory leaks"""
        leaks = []
        gc.collect()
        
        for obj in gc.garbage:
            leaks.append(f"Unreachable: {type(obj).__name__}")
        
        return leaks
    
    def optimize(self) -> Dict:
        """Run memory optimization"""
        before = self.tracker.get_memory_usage()
        
        # Collect garbage
        collected = self.collect()
        
        # Compact memory
        gc.collect()
        
        after = self.tracker.get_memory_usage()
        
        return {
            "objects_collected": collected,
            "memory_freed": before["rss"] - after["rss"],
            "before": before,
            "after": after
        }


class ObjectPool:
    """Object pool for reusing objects"""
    
    def __init__(self):
        self.pools: Dict[str, List] = defaultdict(list)
    
    def acquire(self, obj_type: str, factory: callable) -> Any:
        """Acquire object from pool or create new"""
        if self.pools[obj_type]:
            return self.pools[obj_type].pop()
        return factory()
    
    def release(self, obj_type: str, obj: Any):
        """Release object back to pool"""
        self.pools[obj_type].append(obj)
    
    def clear(self, obj_type: str = None):
        """Clear pool"""
        if obj_type:
            self.pools[obj_type].clear()
        else:
            self.pools.clear()
    
    def get_stats(self) -> Dict:
        """Get pool statistics"""
        return {
            pool_type: len(objects)
            for pool_type, objects in self.pools.items()
        }


# Global memory manager
_memory_optimizer = MemoryOptimizer()
_object_pool = ObjectPool()


def get_memory_optimizer() -> MemoryOptimizer:
    """Get global memory optimizer"""
    return _memory_optimizer


def get_object_pool() -> ObjectPool:
    """Get global object pool"""
    return _object_pool
