"""Streamer patch: stamp ``message.tenant`` and strip it from REST payloads.

PATCHES
-------
* ``opik.message_processing.streamer.Streamer.put``
* ``opik.message_processing.messages.BaseMessage.as_payload_dict``

ROLE
----
``Streamer.put`` is opik's single chokepoint where every outgoing
message enters the background pipeline. Our wrapper hooks it to stamp
``message.tenant`` from the registered ``TenantManager`` resolver
before the message is queued. ``BaseMessage.as_payload_dict`` is the
inverse: it produces the dict that becomes the JSON REST payload, and
we strip ``tenant`` from it so the field never leaks onto the wire.

WHY
---
``message.tenant`` is the bridge between the request task (where the
``PrincipalManager`` ContextVar is valid) and the consumer thread
(where the REST call actually fires). The streamer is the last point
on the request task that sees every message — stamping anywhere later
would lose the principal context. ``as_payload_dict`` is the only
canonical entry point Opik uses to build wire payloads (every
trace/span message subclass calls ``super().as_payload_dict()``), so
patching it is enough to keep ``tenant`` out of every REST call.

V1 SCOPE
--------
Only ``CreateSpan/UpdateSpan/CreateTrace/UpdateTrace`` messages are
stamped. Feedback-score / guardrail / experiment-item handlers are
deliberately NOT wrapped with ``tenant_request_context`` (V2 scope).
If we stamped those types here, the un-patched handlers' existing
``**msg.__dict__`` REST-model construction would forward
``tenant=<TenantContext>`` into pydantic models that use
``extra="allow"`` — silently shipping the field on the wire.

WITHOUT THIS PATCH
------------------
* No tenant stamping → consumer-thread processor patch sees
  ``message.tenant is None`` → ``tenant_request_context(None)`` is a
  no-op → no per-request tenant headers → AIO 403s.
* No payload strip → if ``tenant`` were ever stamped on a message
  whose REST model is strict (``extra="forbid"``), construction
  would raise. Defensive even though current trace/span models accept
  extras.

Idempotent: ``apply()`` only patches once per process.
"""

from __future__ import annotations

import logging

from .multitenancy import TenantManager

LOGGER = logging.getLogger(__name__)

_applied = False


def apply() -> None:
    """Monkey-patch ``Streamer.put`` and ``BaseMessage.as_payload_dict``."""
    global _applied
    if _applied:
        return

    try:
        from opik.message_processing import streamer as _streamer
        from opik.message_processing import messages as _messages
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning(
            "acceldata_aio_tracer multitenancy: opik message_processing "
            "import failed (%s); skipping streamer patch",
            exc,
        )
        return

    Streamer = _streamer.Streamer
    BaseMessage = _messages.BaseMessage

    # The set of message types that should be tenant-stamped at the
    # streamer boundary. Only trace/span single messages — the batched
    # variants (Create*BatchMessage) are built by the batcher patch
    # which already sets tenant on each emitted batch.
    _STAMPABLE_TYPES = (
        _messages.CreateSpanMessage,
        _messages.UpdateSpanMessage,
        _messages.CreateTraceMessage,
        _messages.UpdateTraceMessage,
    )

    _orig_put = Streamer.put

    def _patched_put(self, message):
        if (
            isinstance(message, _STAMPABLE_TYPES)
            and getattr(message, "tenant", None) is None
        ):
            tenant = TenantManager.get_current_tenant()
            if tenant is not None:
                message.tenant = tenant
        return _orig_put(self, message)

    Streamer.put = _patched_put

    _orig_as_payload = BaseMessage.as_payload_dict

    def _patched_as_payload(self):
        data = _orig_as_payload(self)
        # ``tenant`` is set via setattr on stampable message types and
        # rides along inside ``__dict__``. Strip it so REST payloads
        # never include the field.
        data.pop("tenant", None)
        return data

    BaseMessage.as_payload_dict = _patched_as_payload

    _applied = True
