"""OpenTelemetry log shipper for Metaflow MFLOG lines."""

from __future__ import annotations

from opentelemetry._logs.severity import SeverityNumber
from opentelemetry.sdk._logs import LoggerProvider
from opentelemetry.sdk._logs.export import (
    BatchLogRecordProcessor,
    ConsoleLogExporter,
    LogRecordExporter,
)
from opentelemetry.sdk.resources import SERVICE_NAME as OTEL_SERVICE_NAME
from opentelemetry.sdk.resources import Resource

from .attrs import ATTR_LOG_SOURCE, SERVICE_NAME
from .mflog_parser import RUNTIME_SOURCE, MflogLine

_LOGGER_NAME = "metaflow.logs"

TaskAttributes = dict[str, str]


class LogShipper:
    """Ships Metaflow MFLOG lines as OTel :class:`LogRecord` objects.

    Accepts any :class:`~opentelemetry.sdk.logs.LoggerProvider`, making it
    straightforward to inject an ``InMemoryLogExporter``-backed provider in
    tests.

    Args:
        logger_provider: The OTel provider to emit records through.
    """

    def __init__(self, logger_provider: LoggerProvider) -> None:
        self._logger = logger_provider.get_logger(_LOGGER_NAME)

    def ship(self, line: MflogLine, attributes: TaskAttributes) -> None:
        """Emit a single MFLOG line as an OTel :class:`LogRecord`.

        Args:
            line: Parsed MFLOG line to ship.
            attributes: Metaflow task context (flow, run_id, step, …) to
                attach as record attributes alongside the log source.
        """
        is_runtime = line.source == RUNTIME_SOURCE
        self._logger.emit(
            timestamp=int(line.timestamp.timestamp() * 1e9),
            severity_text="DEBUG" if is_runtime else "INFO",
            severity_number=SeverityNumber.DEBUG if is_runtime else SeverityNumber.INFO,
            body=line.message,
            attributes={ATTR_LOG_SOURCE: line.source, **attributes},
        )


def make_logger_provider(
    task_attributes: TaskAttributes,
    exporter: LogRecordExporter | None = None,
) -> LoggerProvider:
    """Create a :class:`~opentelemetry.sdk.logs.LoggerProvider` configured for
    a single Metaflow task.

    The task attributes are embedded in the OTel :class:`~opentelemetry.sdk.resources.Resource`
    so every log record from this provider carries the full task context.

    By default uses an OTLP HTTP exporter; set ``OTEL_EXPORTER_OTLP_ENDPOINT``
    to control the destination.  Falls back to a :class:`ConsoleLogExporter`
    when the OTLP package is not installed.

    Args:
        task_attributes: Metaflow task context keyed by attribute constants
            from :mod:`metaflow_extensions.logs.plugins.attrs`.
        exporter: Optional exporter override (useful for testing).

    Returns:
        A fully configured :class:`~opentelemetry.sdk.logs.LoggerProvider`.
    """
    if exporter is None:
        try:
            from opentelemetry.exporter.otlp.proto.http._log_exporter import (
                OTLPLogExporter,
            )

            exporter = OTLPLogExporter()
        except ImportError:
            exporter = ConsoleLogExporter()

    resource = Resource.create({OTEL_SERVICE_NAME: SERVICE_NAME, **task_attributes})
    provider = LoggerProvider(resource=resource)
    provider.add_log_record_processor(BatchLogRecordProcessor(exporter))
    return provider
