from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Literal, Protocol


TraceLevel = Literal["DEBUG", "INFO", "WARNING", "ERROR"]


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


@dataclass(slots=True)
class TraceContext:
    trace_id: str
    span_id: str
    parent_span_id: str | None = None
    attributes: dict[str, Any] | None = None


class TraceContextFactory(ABC):
    @abstractmethod
    def new_root(self, operation: str) -> TraceContext:
        """Create a new root trace context."""

    @abstractmethod
    def child_of(self, parent: TraceContext, operation: str) -> TraceContext:
        """Create a child trace context."""


@dataclass(frozen=True)
class TraceContextRecord:
    trace_id: str
    alias: str
    parent_id: str | None = None
    type: str | None = None
    event_time: datetime = field(default_factory=utc_now)
    attrs: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class TraceLogMessage:
    trace_id: str
    step: str
    status: str
    message: str
    level: TraceLevel
    event_time: datetime = field(default_factory=utc_now)
    attrs: dict[str, Any] = field(default_factory=dict)


class TraceTransport(Protocol):
    def write_context(self, record: TraceContextRecord) -> None:
        """Persist trace context record."""

    def write_message(self, record: TraceLogMessage) -> None:
        """Persist trace log message."""