"""Lightweight telemetry primitives for EB1 orchestration.

The goal of this module is to isolate structured instrumentation from the
execution code paths so we can observe routing decisions without spreading
logging concerns across the code base.  Telemetry emitters are intentionally
minimal: they collect contextual hints (models, plans, reasoning knobs) and
leave the sinks pluggable.

The default implementation is ``LoggingTelemetry``, which writes JSON payloads
to the ``ember.gateway.eb1.telemetry`` logger at ``INFO`` level for run
summaries and ``DEBUG`` for fine-grained events.  A ``NullTelemetry`` is
provided for environments where logging is undesired.  ``get_telemetry``
returns a singleton recorder; callers should not instantiate telemetry
implementations directly so downstream integrations can override it globally.

This module is self-contained so introducing alternative sinks (metrics,
distributed tracing) remains a small change with predictable blast radius.
"""

from __future__ import annotations

import json
import logging
import threading
import time
import uuid
from dataclasses import dataclass
from typing import Any, Mapping, MutableMapping, Protocol

# ---------------------------------------------------------------------------
# Public telemetry protocol
# ---------------------------------------------------------------------------


class Eb1Telemetry(Protocol):
    """Interface exposed to EB1 orchestration components."""

    def begin_run(self, *, context: Mapping[str, Any]) -> str:
        """Start a telemetry run and return a stable identifier."""

    def record(self, run_id: str, event: str, payload: Mapping[str, Any]) -> None:
        """Emit an event associated with ``run_id``."""

    def end_run(self, run_id: str, *, outcome: Mapping[str, Any]) -> None:
        """Finalize a run and record summary metadata."""


# ---------------------------------------------------------------------------
# Null implementation
# ---------------------------------------------------------------------------


class NullTelemetry:
    """No-op implementation used when telemetry is disabled."""

    def begin_run(self, *, context: Mapping[str, Any]) -> str:  # noqa: D401
        return uuid.uuid4().hex

    def record(self, run_id: str, event: str, payload: Mapping[str, Any]) -> None:  # noqa: D401
        return None

    def end_run(self, run_id: str, *, outcome: Mapping[str, Any]) -> None:  # noqa: D401
        return None


# ---------------------------------------------------------------------------
# Logging implementation
# ---------------------------------------------------------------------------


_TELEMETRY_LOGGER = logging.getLogger("ember.gateway.eb1.telemetry")


def _timestamp_ms() -> int:
    return int(time.time() * 1000)


@dataclass(slots=True)
class _RunRecord:
    run_id: str
    started_ms: int
    context: dict[str, Any]


class LoggingTelemetry:
    """Emit structured events to the EB1 telemetry logger."""

    def __init__(self) -> None:
        self._runs: dict[str, _RunRecord] = {}
        self._lock = threading.Lock()

    def begin_run(self, *, context: Mapping[str, Any]) -> str:
        run_id = uuid.uuid4().hex
        record = _RunRecord(run_id=run_id, started_ms=_timestamp_ms(), context=dict(context))
        with self._lock:
            self._runs[run_id] = record
        if _TELEMETRY_LOGGER.isEnabledFor(logging.DEBUG):
            _emit_log("run_started", run_id, context)
        return run_id

    def record(self, run_id: str, event: str, payload: Mapping[str, Any]) -> None:
        if not _TELEMETRY_LOGGER.isEnabledFor(logging.DEBUG):
            return
        enriched = dict(payload)
        enriched.setdefault("ts_ms", _timestamp_ms())
        _emit_log(event, run_id, enriched)

    def end_run(self, run_id: str, *, outcome: Mapping[str, Any]) -> None:
        record: _RunRecord | None
        with self._lock:
            record = self._runs.pop(run_id, None)
        duration_ms: int | None = None
        if record is not None:
            duration_ms = max(0, _timestamp_ms() - record.started_ms)
        summary: MutableMapping[str, Any] = dict(outcome)
        if duration_ms is not None:
            summary.setdefault("duration_ms", duration_ms)
        if record is not None:
            summary.setdefault("context", record.context)
        _emit_log("run_completed", run_id, summary, level=logging.INFO)


def _emit_log(
    event: str,
    run_id: str,
    payload: Mapping[str, Any],
    *,
    level: int = logging.DEBUG,
) -> None:
    entry = {
        "event": event,
        "run_id": run_id,
        "payload": payload,
    }
    try:
        message = json.dumps(entry, default=str, ensure_ascii=False)
    except TypeError:
        message = json.dumps({"event": event, "run_id": run_id, "error": "serialization_failed"})
    _TELEMETRY_LOGGER.log(level, message)


# ---------------------------------------------------------------------------
# Global accessor
# ---------------------------------------------------------------------------


_TELEMETRY_SINGLETON: Eb1Telemetry | None = None


def get_telemetry() -> Eb1Telemetry:
    """Return the active EB1 telemetry recorder."""

    global _TELEMETRY_SINGLETON
    if _TELEMETRY_SINGLETON is None:
        level = _TELEMETRY_LOGGER.getEffectiveLevel()
        if level <= logging.INFO:
            _TELEMETRY_SINGLETON = LoggingTelemetry()
        else:
            _TELEMETRY_SINGLETON = NullTelemetry()
    return _TELEMETRY_SINGLETON


__all__ = [
    "Eb1Telemetry",
    "NullTelemetry",
    "LoggingTelemetry",
    "get_telemetry",
]
