"""Parser for Metaflow's MFLOG structured log format.

MFLOG lines have the form::

    [MFLOG|0|<UTC timestamp>Z|<source>|<id>]<message>

An optional ``[!]`` prefix marks lines that must be persisted.
Non-MFLOG lines (e.g. raw prints before Metaflow bootstraps) are silently
skipped by :func:`parse_stream`.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from datetime import datetime, timezone

# [!][MFLOG|0|<timestamp>Z|<source>|<id>]<message>
_PATTERN = re.compile(
    rb"(\[!\])?"
    rb"\[MFLOG\|"
    rb"(0)\|"
    rb"(.+?)Z\|"
    rb"(.+?)\|"
    rb"(.+?)\]"
    rb"(.+)",
    re.DOTALL,
)

RUNTIME_SOURCE = "runtime"
TASK_SOURCE = "task"


@dataclass(frozen=True)
class MflogLine:
    """A single parsed MFLOG log line.

    Attributes:
        should_persist: True when the line was prefixed with ``[!]``.
        timestamp: UTC datetime when the line was produced.
        source: Either ``"runtime"`` (Metaflow internals) or ``"task"``
            (user code stdout/stderr).
        id: The Metaflow pathspec, e.g. ``MyFlow/123/train/456``.
        message: The raw log message text.
    """

    should_persist: bool
    timestamp: datetime
    source: str
    id: str
    message: str


def parse_line(raw: bytes) -> MflogLine | None:
    """Parse a single MFLOG-formatted line.

    Args:
        raw: Raw bytes of one log line, with or without a trailing newline.

    Returns:
        A parsed :class:`MflogLine`, or ``None`` if the line does not match
        the MFLOG format.
    """
    m = _PATTERN.match(raw.rstrip(b"\n"))
    if not m:
        return None

    persist_flag, _version, ts_bytes, source_bytes, id_bytes, msg_bytes = m.groups()

    timestamp = datetime.fromisoformat(ts_bytes.decode("utf-8")).replace(tzinfo=timezone.utc)

    return MflogLine(
        should_persist=persist_flag is not None,
        timestamp=timestamp,
        source=source_bytes.decode("utf-8"),
        id=id_bytes.decode("utf-8"),
        message=msg_bytes.decode("utf-8", errors="replace"),
    )


def parse_stream(data: bytes) -> list[MflogLine]:
    """Parse all MFLOG lines from a byte stream, skipping non-MFLOG lines.

    Args:
        data: Raw bytes read from a task log file.

    Returns:
        Parsed :class:`MflogLine` objects in file order.
    """
    result = []
    for raw_line in data.splitlines():
        line = parse_line(raw_line)
        if line is not None:
            result.append(line)
    return result
