"""
Resource monitoring tools for Local Deep Research.

This module provides functionality for tracking CPU, memory and other
system resource usage during the research process.
"""

import threading
import time
from contextlib import contextmanager
from typing import Any, Dict

from loguru import logger

# Try to import psutil, but don't fail if not available
try:
    import psutil

    PSUTIL_AVAILABLE = True
except ImportError:
    PSUTIL_AVAILABLE = False
    logger.warning("psutil not available, resource monitoring will be limited")


class ResourceMonitor:
    """
    Monitor system resource usage during research.

    This class provides methods for tracking CPU, memory, and disk usage
    during system execution. It can be used to identify resource bottlenecks
    and optimize configurations for different hardware environments.
    """

    def __init__(
        self,
        sampling_interval: float = 1.0,
        track_process: bool = True,
        track_system: bool = True,
    ):
        """
        Initialize the resource monitor.

        Args:
            sampling_interval: Seconds between resource usage measurements
            track_process: Whether to track this process's resource usage
            track_system: Whether to track overall system resource usage
        """
        self.sampling_interval = sampling_interval
        self.track_process = track_process
        self.track_system = track_system

        self.monitoring = False
        self.monitor_thread = None

        # Resource usage data
        self.process_data = []
        self.system_data = []
        self.start_time = None
        self.end_time = None

        # Check if we can monitor resources
        self.can_monitor = PSUTIL_AVAILABLE
        if not self.can_monitor:
            logger.warning(
                "Resource monitoring requires psutil. Install with: pip install psutil"
            )

    def start(self):
        """Start monitoring resource usage."""
        if not self.can_monitor:
            logger.warning(
                "Resource monitoring not available (psutil not installed)"
            )
            return

        if self.monitoring:
            logger.warning("Resource monitoring already started")
            return

        self.process_data = []
        self.system_data = []
        self.start_time = time.time()
        self.monitoring = True

        # Start monitoring in a background thread
        self.monitor_thread = threading.Thread(
            target=self._monitor_resources, daemon=True
        )
        self.monitor_thread.start()

        logger.info(
            f"Resource monitoring started with {self.sampling_interval}s interval"
        )

    def stop(self):
        """Stop monitoring resource usage."""
        if not self.monitoring:
            return

        self.monitoring = False
        self.end_time = time.time()

        # Wait for the monitoring thread to finish
        if self.monitor_thread:
            self.monitor_thread.join(timeout=2.0)
            self.monitor_thread = None

        logger.info("Resource monitoring stopped")

    def _monitor_resources(self):
        """Background thread that collects resource usage data."""
        if not PSUTIL_AVAILABLE:
            return

        # Get this process
        current_process = psutil.Process()

        while self.monitoring:
            timestamp = time.time()

            try:
                # Monitor this process
                if self.track_process:
                    process_cpu = current_process.cpu_percent(interval=None)
                    process_memory = current_process.memory_info()

                    self.process_data.append(
                        {
                            "timestamp": timestamp,
                            "cpu_percent": process_cpu,
                            "memory_rss": process_memory.rss,  # Resident Set Size in bytes
                            "memory_vms": process_memory.vms,  # Virtual Memory Size in bytes
                            "memory_shared": getattr(
                                process_memory, "shared", 0
                            ),
                            "num_threads": current_process.num_threads(),
                            "open_files": len(current_process.open_files()),
                            "status": current_process.status(),
                        }
                    )

                # Monitor overall system
                if self.track_system:
                    system_cpu = psutil.cpu_percent(interval=None)
                    system_memory = psutil.virtual_memory()
                    system_disk = psutil.disk_usage("/")

                    self.system_data.append(
                        {
                            "timestamp": timestamp,
                            "cpu_percent": system_cpu,
                            "memory_total": system_memory.total,
                            "memory_available": system_memory.available,
                            "memory_used": system_memory.used,
                            "memory_percent": system_memory.percent,
                            "disk_total": system_disk.total,
                            "disk_used": system_disk.used,
                            "disk_percent": system_disk.percent,
                        }
                    )

            except Exception:
                logger.exception("Error monitoring resources")

            # Sleep until next sampling interval
            time.sleep(self.sampling_interval)

    @contextmanager
    def monitor(self):
        """
        Context manager for monitoring resources during a block of code.

        Example:
            with resource_monitor.monitor():
                # Code to monitor
                do_something_resource_intensive()
        """
        self.start()
        try:
            yield
        finally:
            self.stop()

    def get_process_stats(self) -> Dict[str, Any]:
        """
        Get statistics about this process's resource usage.

        Returns:
            Dictionary with process resource usage statistics
        """
        if not self.process_data:
            return {}

        # Extract data series
        cpu_values = [d["cpu_percent"] for d in self.process_data]
        memory_values = [
            d["memory_rss"] / (1024 * 1024) for d in self.process_data
        ]  # Convert to MB

        # Calculate statistics
        return {
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration": self.end_time - self.start_time
            if self.end_time
            else None,
            "sample_count": len(self.process_data),
            "cpu_min": min(cpu_values) if cpu_values else None,
            "cpu_max": max(cpu_values) if cpu_values else None,
            "cpu_avg": sum(cpu_values) / len(cpu_values)
            if cpu_values
            else None,
            "memory_min_mb": min(memory_values) if memory_values else None,
            "memory_max_mb": max(memory_values) if memory_values else None,
            "memory_avg_mb": (
                sum(memory_values) / len(memory_values)
                if memory_values
                else None
            ),
            "thread_max": (
                max(d["num_threads"] for d in self.process_data)
                if self.process_data
                else None
            ),
        }

    def get_system_stats(self) -> Dict[str, Any]:
        """
        Get statistics about overall system resource usage.

        Returns:
            Dictionary with system resource usage statistics
        """
        if not self.system_data:
            return {}

        # Extract data series
        cpu_values = [d["cpu_percent"] for d in self.system_data]
        memory_values = [d["memory_percent"] for d in self.system_data]
        disk_values = [d["disk_percent"] for d in self.system_data]

        # Calculate statistics
        return {
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration": self.end_time - self.start_time
            if self.end_time
            else None,
            "sample_count": len(self.system_data),
            "cpu_min": min(cpu_values) if cpu_values else None,
            "cpu_max": max(cpu_values) if cpu_values else None,
            "cpu_avg": sum(cpu_values) / len(cpu_values)
            if cpu_values
            else None,
            "memory_min_percent": min(memory_values) if memory_values else None,
            "memory_max_percent": max(memory_values) if memory_values else None,
            "memory_avg_percent": (
                sum(memory_values) / len(memory_values)
                if memory_values
                else None
            ),
            "disk_min_percent": min(disk_values) if disk_values else None,
            "disk_max_percent": max(disk_values) if disk_values else None,
            "disk_avg_percent": (
                sum(disk_values) / len(disk_values) if disk_values else None
            ),
            "memory_total_gb": (
                self.system_data[0]["memory_total"] / (1024**3)
                if self.system_data
                else None
            ),
            "disk_total_gb": (
                self.system_data[0]["disk_total"] / (1024**3)
                if self.system_data
                else None
            ),
        }

    def get_combined_stats(self) -> Dict[str, Any]:
        """
        Get combined resource usage statistics.

        Returns:
            Dictionary with both process and system statistics
        """
        process_stats = self.get_process_stats()
        system_stats = self.get_system_stats()

        # Combine stats
        stats = {
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration": self.end_time - self.start_time
            if self.end_time
            else None,
        }

        # Add process stats with 'process_' prefix
        for key, value in process_stats.items():
            if key not in ["start_time", "end_time", "duration"]:
                stats[f"process_{key}"] = value

        # Add system stats with 'system_' prefix
        for key, value in system_stats.items():
            if key not in ["start_time", "end_time", "duration"]:
                stats[f"system_{key}"] = value

        # Calculate derived metrics
        if (
            process_stats.get("memory_max_mb") is not None
            and system_stats.get("memory_total_gb") is not None
        ):
            # Process memory as percentage of total system memory
            system_memory_mb = system_stats["memory_total_gb"] * 1024
            stats["process_memory_percent"] = (
                (process_stats["memory_max_mb"] / system_memory_mb) * 100
                if system_memory_mb > 0
                else 0
            )

        return stats

    def print_summary(self):
        """Print a formatted summary of resource usage."""
        process_stats = self.get_process_stats()
        system_stats = self.get_system_stats()

        print("\n===== RESOURCE USAGE SUMMARY =====")

        if process_stats:
            print("\n--- Process Resources ---")
            print(
                f"CPU usage: {process_stats.get('cpu_avg', 0):.1f}% avg, "
                f"{process_stats.get('cpu_max', 0):.1f}% peak"
            )
            print(
                f"Memory usage: {process_stats.get('memory_avg_mb', 0):.1f} MB avg, "
                f"{process_stats.get('memory_max_mb', 0):.1f} MB peak"
            )
            print(f"Threads: {process_stats.get('thread_max', 0)} max")

        if system_stats:
            print("\n--- System Resources ---")
            print(
                f"CPU usage: {system_stats.get('cpu_avg', 0):.1f}% avg, "
                f"{system_stats.get('cpu_max', 0):.1f}% peak"
            )
            print(
                f"Memory usage: {system_stats.get('memory_avg_percent', 0):.1f}% avg, "
                f"{system_stats.get('memory_max_percent', 0):.1f}% peak "
                f"(Total: {system_stats.get('memory_total_gb', 0):.1f} GB)"
            )
            print(
                f"Disk usage: {system_stats.get('disk_avg_percent', 0):.1f}% avg "
                f"(Total: {system_stats.get('disk_total_gb', 0):.1f} GB)"
            )

        print("\n===================================")

    def export_data(self) -> Dict[str, Any]:
        """
        Export all collected data.

        Returns:
            Dictionary with all collected resource usage data
        """
        return {
            "start_time": self.start_time,
            "end_time": self.end_time,
            "sampling_interval": self.sampling_interval,
            "process_data": self.process_data,
            "system_data": self.system_data,
        }
