"""
Memory monitoring plugin for Heartbeat.

Collects memory and swap usage statistics using psutil.
"""

import logging
from typing import Dict, Any, Optional

try:
    import psutil
except ImportError:
    psutil = None

from hbd.client.plugin import MonitorPlugin

logger = logging.getLogger(__name__)


class MemoryMonitorPlugin(MonitorPlugin):
    """
    Monitor memory and swap usage.
    
    Collects:
    - Physical memory (RAM) usage and statistics
    - Virtual memory details
    - Swap memory usage and statistics
    - Memory available for applications
    
    Configuration:
        interval: Collection interval in seconds (default: 300)
        include_swap: Include swap statistics (default: True)
    """
    
    name = "memory_monitor"
    interval = 300  # Collect every 5 minutes by default
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Initialize the memory monitor plugin.
        
        Args:
            config: Optional configuration dict with keys:
                   - interval: Collection interval in seconds (default: 300)
                   - include_swap: Include swap statistics (default: True)
        """
        super().__init__(config)
        self.include_swap = self.config.get('include_swap', True)
        self.interval = self.config.get('interval', 300)
        
        if psutil is None:
            raise ImportError("psutil library is required for memory_monitor plugin")
    
    async def initialize(self):
        """Initialize the plugin (check psutil availability)."""
        if psutil is None:
            logger.error("psutil not available - memory_monitor cannot run")
            return False
        
        logger.info(f"Memory monitor initialized (interval: {self.interval}s, swap: {self.include_swap})")
        return True
    
    async def collect(self) -> Dict[str, Any]:
        """
        Collect current memory statistics.
        
        Returns:
            Dictionary with memory metrics:
            - memory_total: Total physical RAM in bytes
            - memory_available: Available memory in bytes
            - memory_used: Used memory in bytes
            - memory_free: Free memory in bytes
            - memory_percent: Memory usage percentage
            - memory_active: Active memory (Unix)
            - memory_inactive: Inactive memory (Unix)
            - memory_buffers: Buffers (Linux)
            - memory_cached: Cached (Linux)
            - memory_shared: Shared (Linux)
            - swap_total: Total swap in bytes (if include_swap)
            - swap_used: Used swap in bytes (if include_swap)
            - swap_free: Free swap in bytes (if include_swap)
            - swap_percent: Swap usage percentage (if include_swap)
            - swap_sin: Bytes swapped in from disk (if include_swap)
            - swap_sout: Bytes swapped out to disk (if include_swap)
        """
        if psutil is None:
            logger.error("psutil not available")
            return {}
        
        try:
            data = await self._collect_metrics()
            logger.debug(f"Collected memory metrics: {len(data)} fields")
            return data
        except Exception as e:
            logger.error(f"Error collecting memory metrics: {e}")
            return {"error": str(e)}
    
    async def _collect_metrics(self) -> Dict[str, Any]:
        """Collect memory metrics from psutil."""
        metrics = {}
        
        # Virtual (physical) memory statistics
        vmem = psutil.virtual_memory()
        metrics['memory_total'] = vmem.total
        metrics['memory_available'] = vmem.available
        metrics['memory_used'] = vmem.used
        metrics['memory_free'] = vmem.free
        metrics['memory_percent'] = vmem.percent
        
        # Platform-specific memory details
        if hasattr(vmem, 'active'):
            metrics['memory_active'] = vmem.active
        if hasattr(vmem, 'inactive'):
            metrics['memory_inactive'] = vmem.inactive
        if hasattr(vmem, 'buffers'):
            metrics['memory_buffers'] = vmem.buffers
        if hasattr(vmem, 'cached'):
            metrics['memory_cached'] = vmem.cached
        if hasattr(vmem, 'shared'):
            metrics['memory_shared'] = vmem.shared
        
        # Swap memory statistics
        if self.include_swap:
            try:
                swap = psutil.swap_memory()
                metrics['swap_total'] = swap.total
                metrics['swap_used'] = swap.used
                metrics['swap_free'] = swap.free
                metrics['swap_percent'] = swap.percent
                
                # Swap in/out counters (may not be available on all platforms)
                if hasattr(swap, 'sin'):
                    metrics['swap_sin'] = swap.sin
                if hasattr(swap, 'sout'):
                    metrics['swap_sout'] = swap.sout
            except Exception as e:
                logger.warning(f"Could not collect swap statistics: {e}")
        
        return metrics
    
    async def cleanup(self):
        """Cleanup (nothing to do for this plugin)."""
        logger.info("Memory monitor cleanup")


# Plugin instance for automatic discovery
plugin = MemoryMonitorPlugin