"""
Resource management utilities for the backtesting system.
Handles memory limits, resource cleanup, and system resource monitoring.
"""

import asyncio
import logging
import os
import resource
import signal
import threading
import time
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Optional

import psutil

from maverick_mcp.utils.memory_profiler import (
    check_memory_leak,
    force_garbage_collection,
    get_memory_stats,
)

logger = logging.getLogger(__name__)

# Resource limits (in bytes)
DEFAULT_MEMORY_LIMIT = 2 * 1024 * 1024 * 1024  # 2GB
DEFAULT_SWAP_LIMIT = 4 * 1024 * 1024 * 1024  # 4GB
CRITICAL_MEMORY_THRESHOLD = 0.9  # 90% of limit

# Global resource manager instance
_resource_manager: Optional["ResourceManager"] = None


@dataclass
class ResourceLimits:
    """Resource limit configuration."""

    memory_limit_bytes: int = DEFAULT_MEMORY_LIMIT
    swap_limit_bytes: int = DEFAULT_SWAP_LIMIT
    cpu_time_limit_seconds: int = 3600  # 1 hour
    file_descriptor_limit: int = 1024
    enable_memory_monitoring: bool = True
    enable_cpu_monitoring: bool = True
    cleanup_interval_seconds: int = 60


@dataclass
class ResourceUsage:
    """Current resource usage snapshot."""

    memory_rss_bytes: int
    memory_vms_bytes: int
    memory_percent: float
    cpu_percent: float
    open_files: int
    threads: int
    timestamp: float


class ResourceExhaustionError(Exception):
    """Raised when resource limits are exceeded."""

    pass


class ResourceManager:
    """System resource manager with limits and cleanup."""

    def __init__(self, limits: ResourceLimits = None):
        """Initialize resource manager.

        Args:
            limits: Resource limits configuration
        """
        self.limits = limits or ResourceLimits()
        self.process = psutil.Process()
        self.monitoring_active = False
        self.cleanup_callbacks: list[Callable[[], None]] = []
        self.resource_history: list[ResourceUsage] = []
        self.max_history_size = 100

        # Setup signal handlers for graceful shutdown
        self._setup_signal_handlers()

        # Start monitoring if enabled
        if self.limits.enable_memory_monitoring or self.limits.enable_cpu_monitoring:
            self.start_monitoring()

    def _setup_signal_handlers(self):
        """Setup signal handlers for resource cleanup."""

        def signal_handler(signum, frame):
            logger.info(f"Received signal {signum}, performing cleanup")
            self.cleanup_all()

        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)

    def start_monitoring(self):
        """Start resource monitoring in background thread."""
        if self.monitoring_active:
            return

        self.monitoring_active = True
        monitor_thread = threading.Thread(target=self._monitor_resources, daemon=True)
        monitor_thread.start()
        logger.info("Resource monitoring started")

    def stop_monitoring(self):
        """Stop resource monitoring."""
        self.monitoring_active = False
        logger.info("Resource monitoring stopped")

    def _monitor_resources(self):
        """Background resource monitoring loop."""
        while self.monitoring_active:
            try:
                usage = self.get_current_usage()
                self.resource_history.append(usage)

                # Keep history size manageable
                if len(self.resource_history) > self.max_history_size:
                    self.resource_history.pop(0)

                # Check limits and trigger cleanup if needed
                self._check_resource_limits(usage)

                time.sleep(self.limits.cleanup_interval_seconds)

            except Exception as e:
                logger.error(f"Error in resource monitoring: {e}")
                time.sleep(30)  # Back off on errors

    def get_current_usage(self) -> ResourceUsage:
        """Get current resource usage."""
        try:
            memory_info = self.process.memory_info()
            cpu_percent = self.process.cpu_percent()

            # Get open files count safely
            try:
                open_files = len(self.process.open_files())
            except (psutil.AccessDenied, psutil.NoSuchProcess):
                open_files = 0

            # Get thread count safely
            try:
                threads = self.process.num_threads()
            except (psutil.AccessDenied, psutil.NoSuchProcess):
                threads = 0

            return ResourceUsage(
                memory_rss_bytes=memory_info.rss,
                memory_vms_bytes=memory_info.vms,
                memory_percent=self.process.memory_percent(),
                cpu_percent=cpu_percent,
                open_files=open_files,
                threads=threads,
                timestamp=time.time(),
            )

        except Exception as e:
            logger.error(f"Error getting resource usage: {e}")
            return ResourceUsage(0, 0, 0, 0, 0, 0, time.time())

    def _check_resource_limits(self, usage: ResourceUsage):
        """Check if resource limits are exceeded and take action."""
        # Memory limit check
        if (
            usage.memory_rss_bytes
            > self.limits.memory_limit_bytes * CRITICAL_MEMORY_THRESHOLD
        ):
            logger.warning(
                f"Memory usage {usage.memory_rss_bytes / (1024**3):.2f}GB "
                f"approaching limit {self.limits.memory_limit_bytes / (1024**3):.2f}GB"
            )
            self._trigger_emergency_cleanup()

        if usage.memory_rss_bytes > self.limits.memory_limit_bytes:
            logger.critical(
                f"Memory limit exceeded: {usage.memory_rss_bytes / (1024**3):.2f}GB "
                f"> {self.limits.memory_limit_bytes / (1024**3):.2f}GB"
            )
            raise ResourceExhaustionError("Memory limit exceeded")

        # File descriptor check
        if usage.open_files > self.limits.file_descriptor_limit * 0.9:
            logger.warning(f"High file descriptor usage: {usage.open_files}")
            self._close_unused_files()

    def _trigger_emergency_cleanup(self):
        """Trigger emergency resource cleanup."""
        logger.info("Triggering emergency resource cleanup")

        # Force garbage collection
        force_garbage_collection()

        # Run cleanup callbacks
        for callback in self.cleanup_callbacks:
            try:
                callback()
            except Exception as e:
                logger.error(f"Error in cleanup callback: {e}")

        # Clear memory profiler snapshots
        try:
            from maverick_mcp.utils.memory_profiler import reset_memory_stats

            reset_memory_stats()
        except ImportError:
            pass

        # Clear cache if available
        try:
            from maverick_mcp.data.cache import clear_cache

            clear_cache()
        except ImportError:
            pass

    def _close_unused_files(self):
        """Close unused file descriptors."""
        try:
            # Get current open files
            open_files = self.process.open_files()
            logger.debug(f"Found {len(open_files)} open files")

            # Note: We can't automatically close files as that might break the application
            # This is mainly for monitoring and alerting
            for file_info in open_files:
                logger.debug(f"Open file: {file_info.path}")

        except Exception as e:
            logger.debug(f"Could not enumerate open files: {e}")

    def add_cleanup_callback(self, callback: Callable[[], None]):
        """Add a cleanup callback function."""
        self.cleanup_callbacks.append(callback)

    def cleanup_all(self):
        """Run all cleanup callbacks and garbage collection."""
        logger.info("Running comprehensive resource cleanup")

        # Run cleanup callbacks
        for callback in self.cleanup_callbacks:
            try:
                callback()
            except Exception as e:
                logger.error(f"Error in cleanup callback: {e}")

        # Force garbage collection
        force_garbage_collection()

        # Log final resource usage
        usage = self.get_current_usage()
        logger.info(
            f"Post-cleanup usage: {usage.memory_rss_bytes / (1024**2):.2f}MB memory, "
            f"{usage.open_files} files, {usage.threads} threads"
        )

    def get_resource_report(self) -> dict[str, Any]:
        """Get comprehensive resource usage report."""
        current = self.get_current_usage()

        report = {
            "current_usage": {
                "memory_mb": current.memory_rss_bytes / (1024**2),
                "memory_percent": current.memory_percent,
                "cpu_percent": current.cpu_percent,
                "open_files": current.open_files,
                "threads": current.threads,
            },
            "limits": {
                "memory_limit_mb": self.limits.memory_limit_bytes / (1024**2),
                "memory_usage_ratio": current.memory_rss_bytes
                / self.limits.memory_limit_bytes,
                "file_descriptor_limit": self.limits.file_descriptor_limit,
            },
            "monitoring": {
                "active": self.monitoring_active,
                "history_size": len(self.resource_history),
                "cleanup_callbacks": len(self.cleanup_callbacks),
            },
        }

        # Add memory profiler stats if available
        try:
            memory_stats = get_memory_stats()
            report["memory_profiler"] = memory_stats
        except Exception:
            pass

        return report

    def set_memory_limit(self, limit_bytes: int):
        """Set memory limit for the process."""
        try:
            # Set soft and hard memory limits
            resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes))
            self.limits.memory_limit_bytes = limit_bytes
            logger.info(f"Memory limit set to {limit_bytes / (1024**3):.2f}GB")
        except Exception as e:
            logger.warning(f"Could not set memory limit: {e}")

    def check_memory_health(self) -> dict[str, Any]:
        """Check memory health and detect potential issues."""
        health_report = {
            "status": "healthy",
            "issues": [],
            "recommendations": [],
        }

        current = self.get_current_usage()

        # Check memory usage
        usage_ratio = current.memory_rss_bytes / self.limits.memory_limit_bytes
        if usage_ratio > 0.9:
            health_report["status"] = "critical"
            health_report["issues"].append(f"Memory usage at {usage_ratio:.1%}")
            health_report["recommendations"].append("Trigger immediate cleanup")
        elif usage_ratio > 0.7:
            health_report["status"] = "warning"
            health_report["issues"].append(f"High memory usage at {usage_ratio:.1%}")
            health_report["recommendations"].append("Consider cleanup")

        # Check for memory leaks
        if check_memory_leak(threshold_mb=100.0):
            health_report["status"] = "warning"
            health_report["issues"].append("Potential memory leak detected")
            health_report["recommendations"].append("Review memory profiler logs")

        # Check file descriptor usage
        fd_ratio = current.open_files / self.limits.file_descriptor_limit
        if fd_ratio > 0.8:
            health_report["status"] = "warning"
            health_report["issues"].append(
                f"High file descriptor usage: {current.open_files}"
            )
            health_report["recommendations"].append("Review open files")

        return health_report


@contextmanager
def resource_limit_context(
    memory_limit_mb: int = None,
    cpu_limit_percent: float = None,
    cleanup_on_exit: bool = True,
):
    """Context manager for resource-limited operations.

    Args:
        memory_limit_mb: Memory limit in MB
        cpu_limit_percent: CPU limit as percentage
        cleanup_on_exit: Whether to cleanup on exit

    Yields:
        ResourceManager instance
    """
    limits = ResourceLimits()
    if memory_limit_mb:
        limits.memory_limit_bytes = memory_limit_mb * 1024 * 1024

    manager = ResourceManager(limits)

    try:
        yield manager
    finally:
        if cleanup_on_exit:
            manager.cleanup_all()
        manager.stop_monitoring()


def get_resource_manager() -> ResourceManager:
    """Get or create global resource manager instance."""
    global _resource_manager
    if _resource_manager is None:
        _resource_manager = ResourceManager()
    return _resource_manager


def set_process_memory_limit(limit_gb: float):
    """Set memory limit for current process.

    Args:
        limit_gb: Memory limit in gigabytes
    """
    limit_bytes = int(limit_gb * 1024 * 1024 * 1024)
    manager = get_resource_manager()
    manager.set_memory_limit(limit_bytes)


def monitor_async_task(task: asyncio.Task, name: str = "unknown"):
    """Monitor an async task for resource usage.

    Args:
        task: Asyncio task to monitor
        name: Name of the task for logging
    """

    def task_done_callback(finished_task):
        if finished_task.exception():
            logger.error(f"Task {name} failed: {finished_task.exception()}")
        else:
            logger.debug(f"Task {name} completed successfully")

        # Trigger cleanup
        manager = get_resource_manager()
        manager._trigger_emergency_cleanup()

    task.add_done_callback(task_done_callback)


class ResourceAwareExecutor:
    """Executor that respects resource limits."""

    def __init__(self, max_workers: int = None, memory_limit_mb: int = None):
        """Initialize resource-aware executor.

        Args:
            max_workers: Maximum worker threads
            memory_limit_mb: Memory limit in MB
        """
        self.max_workers = max_workers or min(32, (os.cpu_count() or 1) + 4)
        self.memory_limit_mb = memory_limit_mb or 500
        self.active_tasks = 0
        self.lock = threading.Lock()

    def submit(self, fn: Callable, *args, **kwargs):
        """Submit a task for execution with resource monitoring."""
        with self.lock:
            if self.active_tasks >= self.max_workers:
                raise ResourceExhaustionError("Too many active tasks")

            # Check memory before starting
            current_usage = get_resource_manager().get_current_usage()
            if current_usage.memory_rss_bytes > self.memory_limit_mb * 1024 * 1024:
                raise ResourceExhaustionError("Memory limit would be exceeded")

            self.active_tasks += 1

        try:
            result = fn(*args, **kwargs)
            return result
        finally:
            with self.lock:
                self.active_tasks -= 1


# Utility functions


def cleanup_on_low_memory(threshold_mb: float = 500.0):
    """Decorator to trigger cleanup when memory is low.

    Args:
        threshold_mb: Memory threshold in MB
    """

    def decorator(func):
        def wrapper(*args, **kwargs):
            get_resource_manager().get_current_usage()
            available_mb = psutil.virtual_memory().available / (1024**2)

            if available_mb < threshold_mb:
                logger.warning(
                    f"Low memory detected ({available_mb:.1f}MB), triggering cleanup"
                )
                get_resource_manager()._trigger_emergency_cleanup()

            return func(*args, **kwargs)

        return wrapper

    return decorator


def log_resource_usage(func: Callable = None, *, interval: int = 60):
    """Decorator to log resource usage periodically.

    Args:
        func: Function to decorate
        interval: Logging interval in seconds
    """

    def decorator(f):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            start_usage = get_resource_manager().get_current_usage()

            try:
                return f(*args, **kwargs)
            finally:
                end_usage = get_resource_manager().get_current_usage()
                duration = time.time() - start_time

                memory_delta = end_usage.memory_rss_bytes - start_usage.memory_rss_bytes
                logger.info(
                    f"{f.__name__} completed in {duration:.2f}s, "
                    f"memory delta: {memory_delta / (1024**2):+.2f}MB"
                )

        return wrapper

    if func is None:
        return decorator
    else:
        return decorator(func)


# Initialize global resource manager
def initialize_resource_management(memory_limit_gb: float = 2.0):
    """Initialize global resource management.

    Args:
        memory_limit_gb: Memory limit in GB
    """
    global _resource_manager

    limits = ResourceLimits(
        memory_limit_bytes=int(memory_limit_gb * 1024 * 1024 * 1024),
        enable_memory_monitoring=True,
        enable_cpu_monitoring=True,
    )

    _resource_manager = ResourceManager(limits)
    logger.info(
        f"Resource management initialized with {memory_limit_gb}GB memory limit"
    )


# Cleanup function for graceful shutdown
def shutdown_resource_management():
    """Shutdown resource management gracefully."""
    global _resource_manager
    if _resource_manager:
        _resource_manager.stop_monitoring()
        _resource_manager.cleanup_all()
        _resource_manager = None
    logger.info("Resource management shut down")
