"""Replay patch: preserve ``message.tenant`` across SQLite round-trip.

PATCHES
-------
``opik.message_processing.messages.from_db_message_dict``

ROLE
----
Opik's connection monitor stores failed messages in an on-disk SQLite
DB and replays them once the server recovers. The store/load goes:

* Write path: ``BaseMessage.as_db_message_dict()`` -> JSON ->
  ``messages`` table.
* Read path: JSON -> dict -> ``from_db_message_dict(MessageClass, data)``
  -> ``BaseMessage`` instance, then re-put through ``Streamer.put``.

Our wrapper attaches ``tenant`` to messages via ``setattr`` (it's not
a real dataclass field on the upstream ``BaseMessage``). The write
path preserves it because ``as_db_message_dict()`` returns
``{**self.__dict__}`` — the runtime attribute rides along. The read
path drops it: ``from_db_message_dict`` only restores keys that appear
in ``dataclasses.fields(message_class)``, and ``tenant`` is not one of
them, so it gets silently filtered out.

This patch wraps ``from_db_message_dict`` to restore ``tenant`` after
the upstream call: read it from the raw ``data`` dict, coerce via
``TenantContext.from_value`` (handles both nested-dict and
already-TenantContext forms), and ``setattr`` the result onto the
deserialized message.

WITHOUT THIS PATCH
------------------
Connection-monitor recovery silently drops tenant identity. After
replay, every queued message has ``tenant=None``, the processor patch
sees no tenant, no headers go on the wire, AIO 403s, and the
connection monitor flips back to "down" — replay fails permanently.
This was the actual bug behind the failing replay test in the
end-to-end run.

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

from __future__ import annotations

import logging

from .multitenancy import TenantContext

LOGGER = logging.getLogger(__name__)

_applied = False


def apply() -> None:
    """Wrap ``from_db_message_dict`` to restore ``tenant``."""
    global _applied
    if _applied:
        return

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

    _orig_from_db = _messages.from_db_message_dict

    def _patched_from_db(message_class, data):
        obj = _orig_from_db(message_class, data)
        raw_tenant = data.get("tenant") if isinstance(data, dict) else None
        if raw_tenant is not None:
            ctx = TenantContext.from_value(raw_tenant)
            if ctx is not None:
                try:
                    obj.tenant = ctx
                except Exception as exc:  # noqa: BLE001
                    LOGGER.debug(
                        "could not restore tenant on %s: %r",
                        type(obj).__name__,
                        exc,
                    )
        return obj

    _messages.from_db_message_dict = _patched_from_db
    _applied = True
