"""Tests for the OTel log shipper."""

from datetime import datetime, timezone

from opentelemetry._logs.severity import SeverityNumber
from opentelemetry.sdk._logs import LoggerProvider
from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter, SimpleLogRecordProcessor

from metaflow_extensions.logs.plugins.attrs import ATTR_LOG_SOURCE
from metaflow_extensions.logs.plugins.exporter import LogShipper
from metaflow_extensions.logs.plugins.mflog_parser import (
    RUNTIME_SOURCE,
    TASK_SOURCE,
    MflogLine,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_TS = datetime(2024, 3, 15, 10, 30, 45, tzinfo=timezone.utc)
_TASK_ATTRS = {
    "metaflow.flow": "MyFlow",
    "metaflow.run_id": "123",
    "metaflow.step": "train",
    "metaflow.task_id": "456",
    "metaflow.attempt": "0",
}


def _make_shipper() -> tuple:
    """Return a (LogShipper, InMemoryLogRecordExporter) pair for testing."""
    exporter = InMemoryLogRecordExporter()
    provider = LoggerProvider()
    provider.add_log_record_processor(SimpleLogRecordProcessor(exporter))
    return LogShipper(provider), exporter


def _task_line(message: str = "hello", source: str = TASK_SOURCE) -> MflogLine:
    return MflogLine(
        should_persist=False,
        timestamp=_TS,
        source=source,
        id="MyFlow/123/train/456",
        message=message,
    )


# ---------------------------------------------------------------------------
# LogShipper
# ---------------------------------------------------------------------------


class TestLogShipper:
    def test_ship_emits_one_record(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(), _TASK_ATTRS)

        assert len(exporter.get_finished_logs()) == 1

    def test_body_is_line_message(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(message="epoch 1 complete"), _TASK_ATTRS)

        record = exporter.get_finished_logs()[0].log_record
        assert record.body == "epoch 1 complete"

    def test_timestamp_matches_mflog_line(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(), _TASK_ATTRS)

        record = exporter.get_finished_logs()[0].log_record
        expected_ns = int(_TS.timestamp() * 1e9)
        assert record.timestamp == expected_ns

    def test_task_source_has_info_severity(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(source=TASK_SOURCE), _TASK_ATTRS)

        record = exporter.get_finished_logs()[0].log_record
        assert record.severity_number == SeverityNumber.INFO
        assert record.severity_text == "INFO"

    def test_runtime_source_has_debug_severity(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(source=RUNTIME_SOURCE), _TASK_ATTRS)

        record = exporter.get_finished_logs()[0].log_record
        assert record.severity_number == SeverityNumber.DEBUG
        assert record.severity_text == "DEBUG"

    def test_log_source_attribute_set_from_mflog_line(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(source=TASK_SOURCE), _TASK_ATTRS)

        attrs = exporter.get_finished_logs()[0].log_record.attributes
        assert attrs is not None
        assert attrs[ATTR_LOG_SOURCE] == TASK_SOURCE

    def test_task_context_attributes_attached(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(), _TASK_ATTRS)

        attrs = exporter.get_finished_logs()[0].log_record.attributes
        assert attrs is not None
        for key, value in _TASK_ATTRS.items():
            assert attrs[key] == value

    def test_ships_multiple_lines_in_order(self):
        shipper, exporter = _make_shipper()

        shipper.ship(_task_line(message="first"), _TASK_ATTRS)
        shipper.ship(_task_line(message="second"), _TASK_ATTRS)
        shipper.ship(_task_line(message="third"), _TASK_ATTRS)

        bodies = [r.log_record.body for r in exporter.get_finished_logs()]
        assert bodies == ["first", "second", "third"]

    def test_different_task_attrs_produce_independent_records(self):
        shipper, exporter = _make_shipper()
        attrs_a = {**_TASK_ATTRS, "metaflow.step": "start"}
        attrs_b = {**_TASK_ATTRS, "metaflow.step": "end"}

        shipper.ship(_task_line(message="msg_a"), attrs_a)
        shipper.ship(_task_line(message="msg_b"), attrs_b)

        records = exporter.get_finished_logs()
        assert records[0].log_record.attributes["metaflow.step"] == "start"
        assert records[1].log_record.attributes["metaflow.step"] == "end"
