# mypy: ignore-errors
# Metaflow's StepDecorator hook methods use untyped dynamic signatures by
# framework design. Type-checking the core logic lives in observer.py instead.
from metaflow.decorators import StepDecorator

from .export.otel import OtelExporter
from .observer import MetricsObserver


class ObservabilityDecorator(StepDecorator):
    """Automatically records step duration, CPU, memory, disk and GPU metrics.

    Usage::

        @observability
        @step
        def train(self):
            ...

    Metrics are exported via OpenTelemetry. Configure the exporter backend
    using standard OTel environment variables (``OTEL_EXPORTER_OTLP_ENDPOINT``,
    etc.) or pass a custom ``MeterProvider`` programmatically.
    """

    name = "observability"
    defaults: dict = {}  # noqa: RUF012 — Metaflow reads this as a class-level dict

    def task_pre_step(
        self,
        step_name,
        task_datastore,
        metadata,
        run_id,
        task_id,
        flow,
        graph,
        retry_count,
        max_user_code_retries,
        ubf_context,
        inputs,
    ) -> None:
        self._observer = MetricsObserver(exporter=OtelExporter())
        self._stored_run_id = run_id
        self._observer.on_step_start(
            step_name=step_name,
            flow_name=flow.__class__.__name__,
            run_id=run_id,
            retry_count=retry_count,
        )

    def task_post_step(
        self,
        step_name,
        flow,
        graph,
        retry_count,
        max_user_code_retries,
    ) -> None:
        self._observer.on_step_end(
            step_name=step_name,
            flow_name=flow.__class__.__name__,
            run_id=self._run_id,
            retry_count=retry_count,
            success=True,
        )

    def task_exception(
        self,
        exception,
        step_name,
        flow,
        graph,
        retry_count,
        max_user_code_retries,
    ) -> None:
        self._observer.on_step_end(
            step_name=step_name,
            flow_name=flow.__class__.__name__,
            run_id=self._run_id,
            retry_count=retry_count,
            success=False,
        )

    @property
    def _run_id(self) -> str:
        return getattr(self, "_stored_run_id", "")
