"""Processor patch: enter ``tenant_request_context`` around each REST call.

PATCHES
-------
``opik.message_processing.processors.online_message_processor.``
``OpikMessageProcessor`` — six trace/span handlers:

* ``_process_create_span_message``
* ``_process_create_trace_message``
* ``_process_update_span_message``
* ``_process_update_trace_message``
* ``_process_create_spans_batch_message``
* ``_process_create_traces_batch_message``

Each is decorator-wrapped with
``tenant_request_context(TenantContext.from_value(message.tenant))``,
which sets a per-thread thread-local for the duration of the underlying
``self._rest_client.{spans|traces}.{create|update|create_*_batch}(...)``
call. ``MultiTenantAuth.auth_flow`` reads that thread-local and
attaches the four ``X-Tenant-*`` / ``X-User-*`` headers to the
outbound httpx Request.

WHY
---
The processor runs on a background ``QueueConsumer`` thread, not the
request thread. By that time the request's ``ContextVar`` is gone, so
we cannot fall back to ``PrincipalManager`` here. Tenant identity has
to ride on the message itself (set by the streamer patch on the
request thread, or the batcher patch when emitting batched messages).
This patch is the bridge from ``message.tenant`` → thread-local →
auth_flow → wire headers.

``TenantContext.from_value`` coerces ``message.tenant`` into a
``TenantContext`` even when it arrives as a plain dict (the shape
that comes back from ``from_db_message_dict`` after a SQLite replay
round-trip — see ``_replay_patch.py`` for why we also patch that).

WITHOUT THIS PATCH
------------------
* No thread-local set → ``MultiTenantAuth.auth_flow`` adds no tenant
  headers → AIO 403s the request.
* Replayed messages whose ``tenant`` is a dict would still have
  routed correctly by the batcher patch's ``from_value`` coercion,
  but the per-request header injection here would have failed.

V1 SCOPE
--------
Only the six trace/span handlers are wrapped. The five non-trace/span
handlers (feedback scores, threads feedback scores, guardrails,
experiment items, attachments) are deferred to V2 — see
``_streamer_patch.py`` for why their messages also stay un-stamped.

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

from __future__ import annotations

import functools
import logging

from .multitenancy import TenantContext
from .multitenant_auth import tenant_request_context

LOGGER = logging.getLogger(__name__)

_applied = False

# Trace/span server-bound handlers that should run inside
# ``tenant_request_context(message.tenant)``.
_HANDLER_NAMES = (
    "_process_create_span_message",
    "_process_create_trace_message",
    "_process_update_span_message",
    "_process_update_trace_message",
    "_process_create_spans_batch_message",
    "_process_create_traces_batch_message",
)


def apply() -> None:
    """Decorator-wrap each trace/span handler. Idempotent."""
    global _applied
    if _applied:
        return

    try:
        from opik.message_processing.processors import (
            online_message_processor as _omp,
        )
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning(
            "acceldata_aio_tracer multitenancy: opik OpikMessageProcessor "
            "import failed (%s); skipping processor patch",
            exc,
        )
        return

    OpikMessageProcessor = _omp.OpikMessageProcessor

    for handler_name in _HANDLER_NAMES:
        original = getattr(OpikMessageProcessor, handler_name, None)
        if original is None:
            LOGGER.warning(
                "acceldata_aio_tracer multitenancy: handler %s missing on "
                "OpikMessageProcessor; skipping wrap",
                handler_name,
            )
            continue

        @functools.wraps(original)
        def _wrapper(self, message, *args, _orig=original, **kwargs):
            # Coerce via TenantContext.from_value: replayed messages
            # come back from SQLite with ``tenant`` as a plain dict, not
            # a TenantContext. tenant_request_context calls
            # ``tenant.headers()`` which only TenantContext implements,
            # so without this coercion replay drains would silently
            # ship without tenant headers and 403.
            ctx = TenantContext.from_value(getattr(message, "tenant", None))
            with tenant_request_context(ctx):
                return _orig(self, message, *args, **kwargs)

        setattr(OpikMessageProcessor, handler_name, _wrapper)

    _applied = True
