"""Metaflow step decorator that ships task logs to an OTel backend."""

from __future__ import annotations

import os
from typing import Any, Callable

from metaflow.decorators import StepDecorator

from .attrs import ATTR_ATTEMPT, ATTR_FLOW, ATTR_RUN_ID, ATTR_STEP, ATTR_TASK_ID
from .exporter import LogShipper, TaskAttributes, make_logger_provider
from .mflog_parser import parse_stream

_LOG_ENV_VARS = ("MFLOG_STDOUT", "MFLOG_STDERR")


class LogsDecorator(StepDecorator):
    """Ships task stdout and stderr to an OpenTelemetry Logs backend.

    Applied per-step::

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

    Configure the export endpoint with standard OTel environment variables
    (``OTEL_EXPORTER_OTLP_ENDPOINT``, etc.).  Log records carry Metaflow
    task context as OTel attributes (flow, run_id, step, task_id, attempt).
    """

    name = "logs"
    defaults: dict[str, Any] = {}

    def __init__(
        self,
        *args: Any,
        provider_factory: Callable[[TaskAttributes], Any] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)  # type: ignore[no-untyped-call]
        self._provider_factory = provider_factory or make_logger_provider
        self._shipper: LogShipper | None = None
        self._provider: Any | None = None
        self._task_attrs: TaskAttributes = {}

    def task_pre_step(
        self,
        step_name: str,
        task_datastore: Any,
        metadata: Any,
        run_id: str,
        task_id: str,
        flow: Any,
        graph: Any,
        retry_count: int,
        max_user_code_retries: int,
        ubf_context: Any,
        inputs: Any,
    ) -> None:
        self._task_attrs = {
            ATTR_FLOW: flow.__class__.__name__,
            ATTR_RUN_ID: str(run_id),
            ATTR_STEP: step_name,
            ATTR_TASK_ID: str(task_id),
            ATTR_ATTEMPT: str(retry_count),
        }
        self._provider = self._provider_factory(self._task_attrs)
        self._shipper = LogShipper(self._provider)

    def task_post_step(
        self,
        step_name: str,
        flow: Any,
        graph: Any,
        retry_count: int,
        max_user_code_retries: int,
    ) -> None:
        self._emit_and_flush()

    def task_exception(
        self,
        exception: Exception,
        step_name: str,
        flow: Any,
        graph: Any,
        retry_count: int,
        max_user_code_retries: int,
    ) -> None:
        self._emit_and_flush()

    def _emit_and_flush(self) -> None:
        if self._shipper is None:
            return
        for env_var in _LOG_ENV_VARS:
            path = os.environ.get(env_var)
            if path and os.path.exists(path):
                with open(path, "rb") as fh:
                    for line in parse_stream(fh.read()):
                        self._shipper.ship(line, self._task_attrs)
        if self._provider is not None:
            self._provider.shutdown()
        self._shipper = None
        self._provider = None
