"""
Memory optimization utilities for large-scale CHRONOS scans.

This module provides strategies for reducing memory usage during
large vulnerability scans, including streaming, chunking, and
efficient data structures.
"""

import gc
import sys
import logging
from typing import Iterator, List, Generator, Optional, Dict, Any, Callable
from dataclasses import dataclass, field
from collections import deque
from contextlib import contextmanager
import tracemalloc
import psutil
import os

logger = logging.getLogger(__name__)


@dataclass
class MemoryStats:
    """Memory usage statistics."""
    current_mb: float = 0.0
    peak_mb: float = 0.0
    available_mb: float = 0.0
    percent_used: float = 0.0
    gc_stats: Dict[str, int] = field(default_factory=dict)


class MemoryMonitor:
    """Monitor and report memory usage."""
    
    def __init__(self):
        self.process = psutil.Process(os.getpid())
        self.peak_memory = 0.0
        self.start_memory = self._get_memory_mb()
        
    def _get_memory_mb(self) -> float:
        """Get current memory usage in MB."""
        return self.process.memory_info().rss / (1024 * 1024)
        
    def get_stats(self) -> MemoryStats:
        """Get current memory statistics."""
        current = self._get_memory_mb()
        self.peak_memory = max(self.peak_memory, current)
        
        vm = self.process.virtual_memory()
        
        return MemoryStats(
            current_mb=current,
            peak_mb=self.peak_memory,
            available_mb=vm.available / (1024 * 1024),
            percent_used=vm.percent,
            gc_stats={
                'collections_0': gc.get_count()[0],
                'collections_1': gc.get_count()[1],
                'collections_2': gc.get_count()[2],
            }
        )
        
    def print_stats(self) -> None:
        """Print memory statistics."""
        stats = self.get_stats()
        logger.info(
            f"Memory: {stats.current_mb:.2f}MB "
            f"(peak: {stats.peak_mb:.2f}MB, "
            f"available: {stats.available_mb:.2f}MB)"
        )


@contextmanager
def memory_tracking(description: str = ""):
    """Context manager to track memory usage.
    
    Usage:
        with memory_tracking("My operation"):
            # Do something memory-intensive
    """
    monitor = MemoryMonitor()
    start_stats = monitor.get_stats()
    
    try:
        yield monitor
    finally:
        end_stats = monitor.get_stats()
        delta = end_stats.current_mb - start_stats.current_mb
        logger.info(
            f"Memory {description}: +{delta:.2f}MB "
            f"({start_stats.current_mb:.2f}MB → {end_stats.current_mb:.2f}MB)"
        )


@contextmanager
def trace_memory_allocations(top_n: int = 20):
    """Trace and report top memory allocations.
    
    Usage:
        with trace_memory_allocations(top_n=30):
            # Memory-intensive operation
    """
    tracemalloc.start()
    try:
        yield
    finally:
        current, peak = tracemalloc.get_traced_memory()
        
        snapshot = tracemalloc.take_snapshot()
        top_stats = snapshot.statistics('lineno')
        
        logger.info(f"Peak memory traced: {peak / 1024 / 1024:.2f}MB")
        logger.info(f"Top {top_n} memory allocations:")
        
        for stat in top_stats[:top_n]:
            logger.info(f"  {stat}")
            
        tracemalloc.stop()


class StreamingVulnerabilityProcessor:
    """Process vulnerabilities one at a time to minimize memory usage."""
    
    def __init__(self, chunk_size: int = 1000):
        self.chunk_size = chunk_size
        
    def process_vulnerabilities(
        self,
        source_data: Iterator[Dict],
        process_func: Callable
    ) -> Generator[Any, None, None]:
        """Process vulnerabilities in chunks.
        
        Args:
            source_data: Iterator of vulnerability dicts
            process_func: Function to apply to each vulnerability
            
        Yields:
            Processed results
        """
        chunk = []
        
        for vuln in source_data:
            chunk.append(vuln)
            
            if len(chunk) >= self.chunk_size:
                # Process chunk
                for item in chunk:
                    result = process_func(item)
                    if result is not None:
                        yield result
                        
                # Clear chunk
                chunk.clear()
                gc.collect()  # Aggressively free memory
                
        # Process remaining items
        for item in chunk:
            result = process_func(item)
            if result is not None:
                yield result


class EfficientDataStructures:
    """Memory-efficient data structures for large datasets."""
    
    @staticmethod
    def create_sparse_dict(initial_data: Dict = None) -> Dict:
        """Create dict with __slots__ for reduced overhead.
        
        This reduces memory per object by 50-80% compared to regular dicts.
        """
        class SparseRecord:
            __slots__ = ['_data']
            
            def __init__(self, data: Dict):
                self._data = data
                
            def __getitem__(self, key):
                return self._data.get(key)
                
            def __setitem__(self, key, value):
                self._data[key] = value
                
        return SparseRecord(initial_data or {})
        
    @staticmethod
    def create_compact_string_pool() -> 'StringPool':
        """Create interned string pool to deduplicate strings."""
        return StringPool()


class StringPool:
    """Deduplicate strings to reduce memory usage."""
    
    def __init__(self):
        self.pool: Dict[str, str] = {}
        
    def intern(self, string: str) -> str:
        """Add string to pool or return existing."""
        if string not in self.pool:
            self.pool[string] = string
            logger.debug(f"Pool size: {len(self.pool)} strings")
        return self.pool[string]
        
    def get_stats(self) -> Dict[str, int]:
        """Get pool statistics."""
        total_size = sum(sys.getsizeof(s) for s in self.pool.values())
        return {
            'unique_strings': len(self.pool),
            'total_size_bytes': total_size,
            'avg_string_bytes': total_size / len(self.pool) if self.pool else 0
        }


class MemoryEfficientVulnerabilityStorage:
    """Store vulnerabilities with minimal memory overhead."""
    
    def __init__(self, max_items: int = 100000):
        self.max_items = max_items
        self.items: deque = deque(maxlen=max_items)
        self.string_pool = StringPool()
        
    def add_vulnerability(self, vuln_dict: Dict) -> None:
        """Add vulnerability, deduplicating strings.
        
        Args:
            vuln_dict: Vulnerability dictionary
        """
        # Intern all string values
        compacted = {}
        for key, value in vuln_dict.items():
            if isinstance(value, str):
                compacted[key] = self.string_pool.intern(value)
            else:
                compacted[key] = value
                
        self.items.append(compacted)
        
    def get_all(self) -> List[Dict]:
        """Get all vulnerabilities."""
        return list(self.items)
        
    def memory_estimate(self) -> Dict[str, float]:
        """Estimate memory usage."""
        item_size = sum(sys.getsizeof(item) for item in self.items)
        pool_size = sum(sys.getsizeof(s) for s in self.string_pool.pool.values())
        
        return {
            'items_bytes': item_size,
            'string_pool_bytes': pool_size,
            'total_bytes': item_size + pool_size,
            'total_mb': (item_size + pool_size) / (1024 * 1024)
        }


class LazyLoadingScanner:
    """Lazily load scan results to avoid loading everything into memory."""
    
    def __init__(self, result_file: str):
        self.result_file = result_file
        self.file_handle = None
        
    def __enter__(self):
        """Open file for lazy reading."""
        self.file_handle = open(self.result_file, 'r')
        return self
        
    def __exit__(self, *args):
        """Close file."""
        if self.file_handle:
            self.file_handle.close()
            
    def read_vulnerabilities(self) -> Generator[Dict, None, None]:
        """Yield vulnerabilities one at a time from file.
        
        Yields:
            Individual vulnerability dicts
        """
        import json
        
        if not self.file_handle:
            raise RuntimeError("Use with context manager")
            
        for line in self.file_handle:
            try:
                vuln = json.loads(line)
                yield vuln
                gc.collect()  # Aggressively free memory
            except json.JSONDecodeError:
                continue


def optimize_for_large_scan(target_size_mb: int = 1000) -> None:
    """Configure CHRONOS for large scans.
    
    Args:
        target_size_mb: Expected scan size in MB
    """
    # Increase GC thresholds to reduce collection frequency
    gc.set_threshold(3000, 10, 10)
    
    # Set memory tracking
    logger.info(f"Optimized for large scans (~{target_size_mb}MB)")
    logger.info("GC thresholds: (3000, 10, 10)")
    
    # Monitor available memory
    vm = psutil.virtual_memory()
    if vm.available / (1024 * 1024) < target_size_mb * 2:
        logger.warning(
            f"Low available memory: {vm.available / (1024 * 1024):.0f}MB "
            f"(need ~{target_size_mb * 2}MB)"
        )


def reset_memory_optimization() -> None:
    """Reset to default memory settings."""
    gc.set_threshold(700, 10, 10)
    logger.info("Memory optimization reset to defaults")


class ChunkedFileProcessor:
    """Process large files in chunks."""
    
    def __init__(self, chunk_size: int = 1024 * 1024):  # 1MB
        self.chunk_size = chunk_size
        
    def process_file(
        self,
        filepath: str,
        process_func: Callable[[str], Any]
    ) -> Generator[Any, None, None]:
        """Process file in chunks.
        
        Args:
            filepath: Path to file
            process_func: Function to process each chunk
            
        Yields:
            Results from processing
        """
        with open(filepath, 'r') as f:
            while True:
                chunk = f.read(self.chunk_size)
                if not chunk:
                    break
                    
                result = process_func(chunk)
                if result is not None:
                    yield result
                    
                gc.collect()


def estimate_scan_memory(
    num_files: int,
    avg_vulnerabilities_per_file: int = 3,
    bytes_per_vuln: int = 1000
) -> Dict[str, float]:
    """Estimate memory needed for a scan.
    
    Args:
        num_files: Number of files to scan
        avg_vulnerabilities_per_file: Average vulns per file
        bytes_per_vuln: Bytes per vulnerability record
        
    Returns:
        Memory estimate breakdown
    """
    total_vulns = num_files * avg_vulnerabilities_per_file
    vuln_memory = total_vulns * bytes_per_vuln
    overhead = vuln_memory * 0.2  # 20% overhead for Python objects
    
    return {
        'vulnerabilities_count': total_vulns,
        'vulnerabilities_mb': vuln_memory / (1024 * 1024),
        'overhead_mb': overhead / (1024 * 1024),
        'total_mb': (vuln_memory + overhead) / (1024 * 1024),
        'recommended_buffer_mb': (vuln_memory + overhead) / (1024 * 1024) * 1.5
    }
