"""Batcher patch: emit one batch per tenant.

PATCHES
-------
* ``opik.message_processing.batching.batchers.CreateSpanMessageBatcher.``
  ``_create_batches_from_accumulated_messages``
* ``opik.message_processing.batching.batchers.CreateTraceMessageBatcher.``
  ``_create_batches_from_accumulated_messages``

ROLE
----
Opik buffers individual ``CreateSpan/CreateTrace`` messages then flushes
them as ``CreateSpansBatchMessage`` / ``CreateTraceBatchMessage`` to
the backend in batches of up to ``max_batch_size`` items. Each batch
message ultimately becomes a single ``POST /v1/private/{spans,traces}/batch``
call. Our patch groups accumulated messages by
``TenantContext.group_key()`` first, then memory-splits within each
group, then emits one outer batch per group — each carrying its own
``tenant`` so the consumer-thread processor handler can route it.

WHY
----
A single HTTP request can carry exactly one tenant's headers. Vanilla
opik's batcher mixes messages from different tenants into one batch,
which would post tenant-A's spans with tenant-B's headers (or vice
versa) and miss-route them. Splitting by tenant in the batcher is the
only place this can be fixed: by the time the consumer thread fires
the REST call, the spans are already packed in one wire payload.

WITHOUT THIS PATCH
------------------
* Multi-tenant traffic intermixed: the first message's tenant would
  win the headers; spans from other tenants in the same batch would
  be miss-routed (silently land in the wrong workspace).
* No assertion failure — single-tenant deployments would still work
  because all messages share ``tenant=None``.

CAVEATS
-------
This is the brittlest of the V1 patches: the method body mostly
mirrors upstream's ``_create_batches_from_accumulated_messages``, so an
opik upgrade that changes the encoding pipeline (``as_payload_dict``
→ ``remove_none_from_dict`` → ``encode_and_anonymize`` →
``SpanWrite``/``TraceWrite``) silently runs stale logic. We pin
``opik==1.11.3`` in pyproject.toml; integration tests catch drift on
upgrade.

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

from __future__ import annotations

import logging
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

LOGGER = logging.getLogger(__name__)

_applied = False


def apply() -> None:
    """Replace the two batcher methods. Idempotent."""
    global _applied
    if _applied:
        return

    try:
        from opik.message_processing import messages as _messages
        from opik.message_processing.batching import (
            batchers as _batchers,
            sequence_splitter as _sequence_splitter,
        )
        from opik.message_processing import encoder_helpers as _encoder_helpers
        from opik import dict_utils as _dict_utils
        from opik.rest_api.types import span_write as _span_write, trace_write as _trace_write
    except Exception as exc:  # noqa: BLE001
        LOGGER.warning(
            "acceldata_aio_tracer multitenancy: opik batcher import failed "
            "(%s); skipping batcher patch",
            exc,
        )
        return

    from .multitenancy import TenantContext

    def _patched_span_create_batches(self):
        groups: Dict[Optional[Tuple[str, str]], List] = defaultdict(list)
        for item in self._accumulated_messages:
            ctx = TenantContext.from_value(getattr(item, "tenant", None))
            key = ctx.group_key() if ctx is not None else None
            groups[key].append(item)

        out: List[_messages.CreateSpansBatchMessage] = []
        for items in groups.values():
            tenant = TenantContext.from_value(getattr(items[0], "tenant", None))
            rest_spans: List[_span_write.SpanWrite] = []
            for item in items:
                span_write_kwargs = item.as_payload_dict()
                cleaned_span_write_kwargs = _dict_utils.remove_none_from_dict(
                    span_write_kwargs
                )
                cleaned_span_write_kwargs = _encoder_helpers.encode_and_anonymize(
                    cleaned_span_write_kwargs,
                    fields_to_anonymize=_messages.CreateSpansBatchMessage.fields_to_anonymize(),
                    object_type="span",
                )
                rest_spans.append(_span_write.SpanWrite(**cleaned_span_write_kwargs))

            memory_limited_batches = _sequence_splitter.split_into_batches(
                items=rest_spans,
                max_payload_size_MB=self._batch_memory_limit_mb,
            )

            for batch in memory_limited_batches:
                m = _messages.CreateSpansBatchMessage(batch=batch)
                if tenant is not None:
                    m.tenant = tenant
                out.append(m)
        return out

    def _patched_trace_create_batches(self):
        groups: Dict[Optional[Tuple[str, str]], List] = defaultdict(list)
        for item in self._accumulated_messages:
            ctx = TenantContext.from_value(getattr(item, "tenant", None))
            key = ctx.group_key() if ctx is not None else None
            groups[key].append(item)

        out: List[_messages.CreateTraceBatchMessage] = []
        for items in groups.values():
            tenant = TenantContext.from_value(getattr(items[0], "tenant", None))
            rest_traces: List[_trace_write.TraceWrite] = []
            for item in items:
                trace_write_kwargs = item.as_payload_dict()
                cleaned_trace_write_kwargs = _dict_utils.remove_none_from_dict(
                    trace_write_kwargs
                )
                cleaned_trace_write_kwargs = _encoder_helpers.encode_and_anonymize(
                    cleaned_trace_write_kwargs,
                    fields_to_anonymize=_messages.CreateTraceBatchMessage.fields_to_anonymize(),
                    object_type="trace",
                )
                rest_traces.append(
                    _trace_write.TraceWrite(**cleaned_trace_write_kwargs)
                )

            memory_limited_batches = _sequence_splitter.split_into_batches(
                items=rest_traces,
                max_payload_size_MB=self._batch_memory_limit_mb,
            )

            for batch in memory_limited_batches:
                m = _messages.CreateTraceBatchMessage(batch=batch)
                if tenant is not None:
                    m.tenant = tenant
                out.append(m)
        return out

    _batchers.CreateSpanMessageBatcher._create_batches_from_accumulated_messages = (
        _patched_span_create_batches
    )
    _batchers.CreateTraceMessageBatcher._create_batches_from_accumulated_messages = (
        _patched_trace_create_batches
    )
    _applied = True
