import psutil
import threading
from ..ctx.common import AbstractBothContextManager as AbstractBothContextManager
from ..print.message import debug as debug, info as info, warning as warning
from typing import Any

class ProcessMetricsMonitor(AbstractBothContextManager['ProcessMetricsMonitor']):
    ''' Monitor that collects CPU, memory, I/O, and thread metrics for a specific
\tprocess and (optionally) all its children, then logs them to MLflow.

\tThis is the per-process counterpart of MLflow\'s built-in
\t``log_system_metrics=True`` which only captures **system-wide** metrics.
\tHere every metric is scoped to the process tree rooted at *pid*.

\tMetrics collected (all prefixed with ``process/``):

\t- ``cpu_usage_percentage`` - cumulative CPU % (sum over the tree)
\t- ``memory_rss_megabytes`` - resident set size in MB
\t- ``memory_vms_megabytes`` - virtual memory size in MB
\t- ``memory_uss_megabytes`` - unique set size in MB (Linux only, falls back to RSS)
\t- ``memory_usage_percentage`` - RSS as % of total available RAM (see *max_memory_megabytes*)
\t- ``num_threads`` - total thread count across the tree
\t- ``num_fds`` - total open file descriptors (Linux only, 0 on other OS)
\t- ``io_read_megabytes`` - cumulative bytes read in MB (since process start)
\t- ``io_write_megabytes`` - cumulative bytes written in MB (since process start)

\tArgs:
\t\tpid                     (int):      PID of the root process to monitor. Defaults to the current process (``os.getpid()``).
\t\tchildren                (bool):     Whether to include child processes (recursively) in the metrics. Defaults to True.
\t\tsampling_interval       (float):    Seconds between each sample collection. Defaults to 10.
\t\tsamples_before_logging  (int):      Number of samples to average before logging. Defaults to 1.
\t\tprefix                  (str):      Metric name prefix. Defaults to ``"process/"``.
\t\tverbose                 (bool):     Whether to log verbose debug messages. Defaults to False.
\t\tmax_memory_megabytes    (float):    Override the total memory in MB used to compute ``memory_usage_percentage``.
\t\t\tUseful in containerized environments (e.g. Kubernetes pods) where ``psutil`` reports the
\t\t\thost\'s total RAM instead of the container\'s limit. Defaults to ``None`` (use system total).
\t\tmax_cpu_count           (float):    Override the number of CPUs used to normalise ``cpu_usage_percentage``.
\t\t\tFor example, set to ``8.0`` when a pod is limited to 8 cores on a 128-core host.
\t\t\tDefaults to ``None`` (use ``os.cpu_count()``).

\tExamples:
\t\t.. code-block:: python

\t\t\t> import mlflow
\t\t\t> from stouputils.mlflow.process_metrics_monitor import ProcessMetricsMonitor
\t\t\t> mlflow.set_experiment("my_experiment")
\t\t\t> with mlflow.start_run():
\t\t\t.     monitor = ProcessMetricsMonitor(pid=12345, children=True, sampling_interval=5)
\t\t\t.     monitor.start()
\t\t\t.     # ... do heavy work ...
\t\t\t.     monitor.finish()

\t\tOr as a context manager:

\t\t.. code-block:: python

\t\t\t> import mlflow
\t\t\t> from stouputils.mlflow.process_metrics_monitor import ProcessMetricsMonitor
\t\t\t> mlflow.set_experiment("my_experiment")
\t\t\t> with mlflow.start_run():
\t\t\t.     with ProcessMetricsMonitor(pid=12345):
\t\t\t.         # ... do heavy work ...
\t\t\t.         pass
\t'''
    pid: int
    children: bool
    sampling_interval: float
    samples_before_logging: int
    prefix: str
    verbose: bool
    max_memory_megabytes: float
    max_cpu_count: float
    run_id: str | None
    shutdown_event: threading.Event
    thread: threading.Thread | None
    step: int
    samples: list[dict[str, float]]
    processes: dict[int, psutil.Process]
    def __init__(self, pid: int | None = None, children: bool = True, sampling_interval: float = 10.0, samples_before_logging: int = 1, prefix: str = 'process/', verbose: bool = False, max_memory_megabytes: float | None = None, max_cpu_count: float | None = None) -> None: ...
    def __enter__(self) -> ProcessMetricsMonitor: ...
    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None: ...
    async def __aenter__(self) -> ProcessMetricsMonitor: ...
    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None: ...
    def start(self) -> None:
        """ Start the background monitoring thread. """
    def finish(self) -> None:
        """ Stop monitoring and flush remaining metrics to MLflow. """
    def collect_once(self) -> dict[str, float]:
        """ Collect one snapshot of metrics for the process tree.

\t\tReturns:
\t\t\tdict[str, float]:\tA dictionary of metric names to values.
\t\t"""
    def aggregate(self, samples: list[dict[str, float]]) -> dict[str, float]:
        """ Average the collected samples.

\t\tArgs:
\t\t\tsamples (list[dict[str, float]]):\tList of metric dictionaries.

\t\tReturns:
\t\t\tdict[str, float]:\tA dictionary of averaged metric values.
\t\t"""
    def publish(self, metrics: dict[str, float]) -> None:
        """ Log the aggregated metrics to the active MLflow run.

\t\tArgs:
\t\t\tmetrics (dict[str, float]):\tAggregated metric values.
\t\t"""
    def monitor_loop(self) -> None:
        """ Main monitoring loop running in a daemon thread. """
    def flush_remaining(self) -> None:
        """ Flush any buffered samples that haven't been logged yet. """
