"""Enhanced metrics collection for task tracking.

This module provides comprehensive performance metrics collection during node execution,
including CPU usage, memory consumption, I/O operations, and custom metrics.

Design Purpose:
- Enable real-time performance monitoring during node execution
- Integrate seamlessly with TaskManager and visualization components
- Support both synchronous and asynchronous execution contexts

Upstream Dependencies:
- Runtime engines (local.py, parallel.py, docker.py) use this to collect metrics
- TaskManager uses this to store performance data

Downstream Consumers:
- Visualization components use collected metrics for performance graphs
- Export utilities include metrics in workflow reports
"""

import asyncio
import threading
import time
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any

try:
    import psutil

    PSUTIL_AVAILABLE = True
except ImportError:
    psutil = None  # type: ignore[assignment]
    PSUTIL_AVAILABLE = False


@dataclass
class PerformanceMetrics:
    """Container for comprehensive performance metrics.

    Attributes:
        duration: Execution time in seconds
        cpu_percent: Average CPU usage percentage
        memory_mb: Peak memory usage in MB
        memory_delta_mb: Memory increase during execution
        io_read_bytes: Total bytes read during execution
        io_write_bytes: Total bytes written during execution
        io_read_count: Number of read operations
        io_write_count: Number of write operations
        thread_count: Number of threads used
        context_switches: Number of context switches
        custom: Dictionary of custom metrics
    """

    duration: float = 0.0
    cpu_percent: float = 0.0
    memory_mb: float = 0.0
    memory_delta_mb: float = 0.0
    io_read_bytes: int = 0
    io_write_bytes: int = 0
    io_read_count: int = 0
    io_write_count: int = 0
    thread_count: int = 1
    context_switches: int = 0
    custom: dict[str, Any] = field(default_factory=dict)

    def to_task_metrics(self) -> dict[str, Any]:
        """Convert to TaskMetrics compatible format."""
        return {
            "duration": self.duration,
            "memory_usage_mb": self.memory_mb,
            "cpu_usage": self.cpu_percent,
            "custom_metrics": {
                "memory_delta_mb": self.memory_delta_mb,
                "io_read_bytes": self.io_read_bytes,
                "io_write_bytes": self.io_write_bytes,
                "io_read_count": self.io_read_count,
                "io_write_count": self.io_write_count,
                "thread_count": self.thread_count,
                "context_switches": self.context_switches,
                **self.custom,
            },
        }


class MetricsCollector:
    """Collects performance metrics during task execution.

    This class provides context managers for collecting detailed performance
    metrics during node execution, with support for both process-level and
    system-level monitoring.

    Usage:
        >>> collector = MetricsCollector()
        >>> with collector.collect() as metrics:
        ...     # Execute node code here
        ...     pass
        >>> performance_data = metrics.result()
    """

    def __init__(
        self,
        sampling_interval: float = 0.1,
        enable_resource_monitoring: bool = True,
    ):
        """Initialize metrics collector.

        Args:
            sampling_interval: How often to sample metrics (seconds)
            enable_resource_monitoring: Whether to enable psutil-based resource monitoring.
                When False, only duration is tracked (no background threads spawned).
                P0D-001: Dramatically reduces per-node overhead by avoiding thread
                creation/join when resource monitoring is not needed.
        """
        self.sampling_interval = sampling_interval
        # P0D-001: Only enable psutil monitoring when both available AND requested
        self._monitoring_enabled = PSUTIL_AVAILABLE and enable_resource_monitoring

        if not PSUTIL_AVAILABLE and enable_resource_monitoring:
            import warnings

            warnings.warn(
                "psutil not available. Performance metrics will be limited to duration only. "
                "Install psutil for comprehensive metrics: pip install psutil"
            )

    @contextmanager
    def collect(self, node_id: str | None = None):
        """Context manager for collecting metrics during execution.

        Args:
            node_id: Optional node identifier for tracking

        Yields:
            MetricsContext: Context object with result() method
        """
        context = MetricsContext(
            node_id=node_id,
            sampling_interval=self.sampling_interval,
            monitoring_enabled=self._monitoring_enabled,
        )

        try:
            context.start()
            yield context
        finally:
            context.stop()

    async def collect_async(self, coro, node_id: str | None = None):
        """Collect metrics for async execution.

        Args:
            coro: Coroutine to execute
            node_id: Optional node identifier

        Returns:
            Tuple of (result, metrics)
        """
        context = MetricsContext(
            node_id=node_id,
            sampling_interval=self.sampling_interval,
            monitoring_enabled=self._monitoring_enabled,
        )

        try:
            context.start()
            result = await coro
            return result, context.result()
        finally:
            context.stop()


class MetricsContext:
    """Context for collecting metrics during a specific execution."""

    def __init__(
        self, node_id: str | None, sampling_interval: float, monitoring_enabled: bool
    ):
        self.node_id = node_id
        self.sampling_interval = sampling_interval
        self.monitoring_enabled = monitoring_enabled

        self.start_time: float | None = None
        self.end_time: float | None = None
        self.process: Any | None = None
        self.initial_io: Any | None = None
        self.initial_memory: float | None = None
        self.peak_memory: float = 0.0
        self.cpu_samples: list = []
        self.monitoring_thread: threading.Thread | None = None
        self._stop_monitoring = threading.Event()

    def start(self):
        """Start metrics collection."""
        self.start_time = time.time()

        if self.monitoring_enabled and psutil is not None:
            try:
                self.process = psutil.Process()
                assert self.process is not None
                mem_info = self.process.memory_info()
                self.initial_memory = mem_info.rss / 1024 / 1024  # MB
                self.peak_memory = (
                    self.initial_memory if self.initial_memory is not None else 0.0
                )

                # Get initial I/O counters if available
                if hasattr(self.process, "io_counters"):
                    try:
                        self.initial_io = self.process.io_counters()
                    except (psutil.AccessDenied, AttributeError):
                        self.initial_io = None

                # Start monitoring thread
                self._stop_monitoring.clear()
                self.monitoring_thread = threading.Thread(
                    target=self._monitor_resources
                )
                self.monitoring_thread.daemon = True
                self.monitoring_thread.start()

            except Exception:
                # Fallback if process monitoring fails
                self.monitoring_enabled = False

    def stop(self):
        """Stop metrics collection."""
        self.end_time = time.time()

        if self.monitoring_enabled and self.monitoring_thread:
            self._stop_monitoring.set()
            self.monitoring_thread.join(timeout=1.0)

    def _monitor_resources(self):
        """Monitor resources in background thread."""
        if self.process is None:
            return
        while not self._stop_monitoring.is_set():
            try:
                # Sample CPU usage
                cpu = self.process.cpu_percent(interval=None)
                if cpu > 0:  # Filter out initial 0 readings
                    self.cpu_samples.append(cpu)

                # Track peak memory
                mem_info = self.process.memory_info()
                memory = mem_info.rss / 1024 / 1024  # MB
                self.peak_memory = max(self.peak_memory, memory)

            except Exception:
                break

            self._stop_monitoring.wait(self.sampling_interval)

    def result(self) -> PerformanceMetrics:
        """Get collected metrics."""
        metrics = PerformanceMetrics()

        # Calculate duration
        if self.start_time and self.end_time:
            metrics.duration = self.end_time - self.start_time

        if self.monitoring_enabled and self.process and psutil is not None:
            try:
                # CPU usage (average of samples)
                if self.cpu_samples:
                    metrics.cpu_percent = sum(self.cpu_samples) / len(self.cpu_samples)

                # Memory metrics
                metrics.memory_mb = self.peak_memory
                if self.initial_memory:
                    current_memory = self.process.memory_info().rss / 1024 / 1024
                    metrics.memory_delta_mb = current_memory - self.initial_memory

                # I/O metrics
                if self.initial_io and hasattr(self.process, "io_counters"):
                    try:
                        current_io = self.process.io_counters()
                        metrics.io_read_bytes = (
                            current_io.read_bytes - self.initial_io.read_bytes
                        )
                        metrics.io_write_bytes = (
                            current_io.write_bytes - self.initial_io.write_bytes
                        )
                        metrics.io_read_count = (
                            current_io.read_count - self.initial_io.read_count
                        )
                        metrics.io_write_count = (
                            current_io.write_count - self.initial_io.write_count
                        )
                    except (psutil.AccessDenied, AttributeError):
                        pass

                # Thread and context switch info
                try:
                    metrics.thread_count = self.process.num_threads()
                    if hasattr(self.process, "num_ctx_switches"):
                        ctx = self.process.num_ctx_switches()
                        metrics.context_switches = ctx.voluntary + ctx.involuntary
                except (psutil.AccessDenied, AttributeError):
                    pass

            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass

        return metrics

    def add_custom_metric(self, name: str, value: Any):
        """Add a custom metric."""
        if not hasattr(self, "_custom_metrics"):
            self._custom_metrics = {}
        self._custom_metrics[name] = value

    def get_custom_metrics(self) -> dict[str, Any]:
        """Get custom metrics."""
        return getattr(self, "_custom_metrics", {})


# Global collector instance for convenience
default_collector = MetricsCollector()


def collect_metrics(func: Callable | None = None, *, node_id: str | None = None):
    """Decorator for collecting metrics on function execution.

    Can be used as @collect_metrics or @collect_metrics(node_id="my_node")

    Args:
        func: Function to wrap
        node_id: Optional node identifier

    Returns:
        Wrapped function that returns (result, metrics) tuple
    """

    def decorator(f):
        if asyncio.iscoroutinefunction(f):

            async def async_wrapper(*args, **kwargs):
                result, metrics = await default_collector.collect_async(
                    f(*args, **kwargs), node_id=node_id
                )
                return result, metrics

            return async_wrapper
        else:

            def sync_wrapper(*args, **kwargs):
                with default_collector.collect(node_id=node_id) as context:
                    result = f(*args, **kwargs)
                return result, context.result()

            return sync_wrapper

    if func is None:
        # Called with arguments: @collect_metrics(node_id="...")
        return decorator
    else:
        # Called without arguments: @collect_metrics
        return decorator(func)
