"""Tests for the LogsDecorator step decorator."""

import os
from unittest.mock import MagicMock, patch

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

from metaflow_extensions.logs.plugins.attrs import (
    ATTR_ATTEMPT,
    ATTR_FLOW,
    ATTR_RUN_ID,
    ATTR_STEP,
    ATTR_TASK_ID,
)
from metaflow_extensions.logs.plugins.decorator import LogsDecorator
from metaflow_extensions.logs.plugins.mflog_parser import TASK_SOURCE

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

_TS = "2024-03-15T10:30:45.000000"
_PATHSPEC = "MyFlow/123/train/456"


def _mflog_line(message: str) -> bytes:
    return f"[MFLOG|0|{_TS}Z|{TASK_SOURCE}|{_PATHSPEC}]{message}".encode()


def _make_flow(name: str = "MyFlow") -> MagicMock:
    flow = MagicMock()
    flow.__class__.__name__ = name
    return flow


def _make_provider_factory() -> tuple:
    """Return (factory_fn, InMemoryLogRecordExporter) for injecting into the decorator."""
    exporter = InMemoryLogRecordExporter()

    def factory(task_attrs: dict) -> LoggerProvider:
        provider = LoggerProvider()
        provider.add_log_record_processor(SimpleLogRecordProcessor(exporter))
        return provider

    return factory, exporter


def _pre_step(
    decorator: LogsDecorator,
    step_name: str = "train",
    run_id: str = "123",
    task_id: str = "456",
    retry_count: int = 0,
    flow_name: str = "MyFlow",
) -> None:
    decorator.task_pre_step(
        step_name=step_name,
        task_datastore=MagicMock(),
        metadata=MagicMock(),
        run_id=run_id,
        task_id=task_id,
        flow=_make_flow(flow_name),
        graph=MagicMock(),
        retry_count=retry_count,
        max_user_code_retries=3,
        ubf_context=None,
        inputs=None,
    )


def _post_step(decorator: LogsDecorator, step_name: str = "train") -> None:
    decorator.task_post_step(
        step_name=step_name,
        flow=_make_flow(),
        graph=MagicMock(),
        retry_count=0,
        max_user_code_retries=3,
    )


def _exception_step(decorator: LogsDecorator) -> None:
    decorator.task_exception(
        exception=RuntimeError("boom"),
        step_name="train",
        flow=_make_flow(),
        graph=MagicMock(),
        retry_count=0,
        max_user_code_retries=3,
    )


# ---------------------------------------------------------------------------
# task_pre_step — builds task context
# ---------------------------------------------------------------------------


class TestPreStep:
    def test_task_attrs_contain_flow_name(self):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec, flow_name="TrainingFlow")

        assert dec._task_attrs[ATTR_FLOW] == "TrainingFlow"

    def test_task_attrs_contain_run_id(self):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec, run_id="42")

        assert dec._task_attrs[ATTR_RUN_ID] == "42"

    def test_task_attrs_contain_step_name(self):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec, step_name="process")

        assert dec._task_attrs[ATTR_STEP] == "process"

    def test_task_attrs_contain_task_id(self):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec, task_id="789")

        assert dec._task_attrs[ATTR_TASK_ID] == "789"

    def test_task_attrs_contain_attempt(self):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec, retry_count=2)

        assert dec._task_attrs[ATTR_ATTEMPT] == "2"

    def test_factory_called_with_task_attrs(self):
        factory = MagicMock(return_value=MagicMock())
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec)

        factory.assert_called_once()
        passed_attrs = factory.call_args[0][0]
        assert ATTR_FLOW in passed_attrs


# ---------------------------------------------------------------------------
# task_post_step — ships logs from MFLOG files
# ---------------------------------------------------------------------------


class TestPostStep:
    def test_ships_stdout_lines(self, tmp_path):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)
        log_file = tmp_path / "stdout.log"
        log_file.write_bytes(b"\n".join([_mflog_line("line one"), _mflog_line("line two")]))

        _pre_step(dec)
        with patch.dict(os.environ, {"MFLOG_STDOUT": str(log_file), "MFLOG_STDERR": ""}):
            _post_step(dec)

        bodies = [r.log_record.body for r in exporter.get_finished_logs()]
        assert bodies == ["line one", "line two"]

    def test_ships_stderr_lines(self, tmp_path):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)
        log_file = tmp_path / "stderr.log"
        log_file.write_bytes(_mflog_line("error output"))

        _pre_step(dec)
        with patch.dict(os.environ, {"MFLOG_STDOUT": "", "MFLOG_STDERR": str(log_file)}):
            _post_step(dec)

        bodies = [r.log_record.body for r in exporter.get_finished_logs()]
        assert "error output" in bodies

    def test_skips_missing_stdout_file(self, tmp_path):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec)
        with patch.dict(
            os.environ,
            {"MFLOG_STDOUT": str(tmp_path / "nonexistent.log"), "MFLOG_STDERR": ""},
        ):
            _post_step(dec)  # must not raise

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

    def test_skips_when_env_var_not_set(self):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec)
        with patch.dict(os.environ, {}, clear=True):
            _post_step(dec)

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

    def test_task_attrs_attached_to_shipped_records(self, tmp_path):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)
        log_file = tmp_path / "stdout.log"
        log_file.write_bytes(_mflog_line("msg"))

        _pre_step(dec, step_name="encode", run_id="99")
        with patch.dict(os.environ, {"MFLOG_STDOUT": str(log_file), "MFLOG_STDERR": ""}):
            _post_step(dec, step_name="encode")

        attrs = exporter.get_finished_logs()[0].log_record.attributes
        assert attrs is not None
        assert attrs[ATTR_STEP] == "encode"
        assert attrs[ATTR_RUN_ID] == "99"

    def test_shipper_is_cleared_after_post_step(self, tmp_path):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec)
        with patch.dict(os.environ, {"MFLOG_STDOUT": "", "MFLOG_STDERR": ""}):
            _post_step(dec)

        assert dec._shipper is None

    def test_noop_when_pre_step_not_called(self):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _post_step(dec)  # must not raise

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


# ---------------------------------------------------------------------------
# task_exception — ships logs on failure
# ---------------------------------------------------------------------------


class TestExceptionStep:
    def test_ships_logs_on_exception(self, tmp_path):
        factory, exporter = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)
        log_file = tmp_path / "stdout.log"
        log_file.write_bytes(_mflog_line("before crash"))

        _pre_step(dec)
        with patch.dict(os.environ, {"MFLOG_STDOUT": str(log_file), "MFLOG_STDERR": ""}):
            _exception_step(dec)

        bodies = [r.log_record.body for r in exporter.get_finished_logs()]
        assert "before crash" in bodies

    def test_shipper_is_cleared_after_exception(self, tmp_path):
        factory, _ = _make_provider_factory()
        dec = LogsDecorator(provider_factory=factory)

        _pre_step(dec)
        with patch.dict(os.environ, {"MFLOG_STDOUT": "", "MFLOG_STDERR": ""}):
            _exception_step(dec)

        assert dec._shipper is None
