import time
from typing import Optional

from .export.protocol import MetricsExporter
from .samplers.cpu import CpuMemorySampler
from .samplers.disk import DiskSampler
from .samplers.gpu import GpuSampler


class MetricsObserver:
    """Coordinates samplers and exports collected metrics via an exporter.

    This class contains the core observability logic and has no dependency on
    Metaflow, making it independently testable. The ``ObservabilityDecorator``
    is a thin Metaflow adapter around this class.
    """

    def __init__(
        self,
        exporter: MetricsExporter,
        cpu_sampler: Optional[CpuMemorySampler] = None,
        disk_sampler: Optional[DiskSampler] = None,
        gpu_sampler: Optional[GpuSampler] = None,
    ) -> None:
        self._exporter = exporter
        self._cpu_sampler = cpu_sampler or CpuMemorySampler()
        self._disk_sampler = disk_sampler or DiskSampler()
        self._gpu_sampler = gpu_sampler or GpuSampler()
        self._start_time: Optional[float] = None

    def on_step_start(
        self,
        step_name: str,
        flow_name: str,
        run_id: str,
        retry_count: int,
    ) -> None:
        self._start_time = time.monotonic()
        self._cpu_sampler.start()
        self._disk_sampler.start()
        self._gpu_sampler.start()

    def on_step_end(
        self,
        step_name: str,
        flow_name: str,
        run_id: str,
        retry_count: int,
        success: bool,
    ) -> None:
        duration = time.monotonic() - (self._start_time or 0.0)

        cpu_stats = self._cpu_sampler.stop()
        disk_stats = self._disk_sampler.stop()
        gpu_stats = self._gpu_sampler.stop()

        attrs = {"step_name": step_name, "flow_name": flow_name, "run_id": run_id}

        self._exporter.record_duration("step.duration", duration, attrs)

        self._exporter.record_gauge("step.cpu.avg_percent", cpu_stats.avg_cpu_percent, attrs)
        self._exporter.record_gauge("step.cpu.max_percent", cpu_stats.max_cpu_percent, attrs)
        self._exporter.record_gauge("step.cpu.p95_percent", cpu_stats.p95_cpu_percent, attrs)
        self._exporter.record_gauge("step.memory.avg_mb", cpu_stats.avg_memory_mb, attrs)
        self._exporter.record_gauge("step.memory.max_mb", cpu_stats.max_memory_mb, attrs)

        self._exporter.record_counter("step.disk.read_bytes", disk_stats.read_bytes, attrs)
        self._exporter.record_counter("step.disk.write_bytes", disk_stats.write_bytes, attrs)
        self._exporter.record_gauge(
            "step.disk.read_throughput_mbps", disk_stats.avg_read_throughput_mbps, attrs
        )
        self._exporter.record_gauge(
            "step.disk.write_throughput_mbps", disk_stats.avg_write_throughput_mbps, attrs
        )

        for device in gpu_stats.devices:
            device_attrs = {**attrs, "gpu_index": str(device.device_index)}
            self._exporter.record_gauge(
                "step.gpu.avg_utilization_percent",
                device.avg_utilization_percent,
                device_attrs,
            )
            self._exporter.record_gauge(
                "step.gpu.max_utilization_percent",
                device.max_utilization_percent,
                device_attrs,
            )
            self._exporter.record_gauge(
                "step.gpu.avg_memory_used_mb",
                device.avg_memory_used_mb,
                device_attrs,
            )
            self._exporter.record_gauge(
                "step.gpu.max_memory_used_mb",
                device.max_memory_used_mb,
                device_attrs,
            )

        self._exporter.record_gauge("step.retry_count", float(retry_count), attrs)

        if not success:
            self._exporter.record_counter("step.failures", 1, attrs)
