"""Multi-tenant routing primitives for the Acceldata AIO tracer.

Apps that route traces to per-tenant workspaces register a resolver via
``TenantManager.set_resolver(...)``. The wrapper's runtime patches call
the resolver inside ``Streamer.put`` (on the request thread/coroutine,
while the app's ``ContextVar`` is still valid) and stamp the resulting
``TenantContext`` on each outbound trace/span message. The processor
patches then read that tenant on the consumer thread and inject the
four ``X-Tenant-*`` / ``X-User-*`` headers into the outgoing HTTP
request.

When no resolver is registered, ``get_current_tenant()`` returns
``None`` and the wrapper behaves exactly like single-tenant mode (no
tenant headers, single-batch flush, identical to vanilla Opik).
"""

from __future__ import annotations

import dataclasses
import logging
from typing import Any, Callable, Dict, Optional, Tuple, Union

LOGGER = logging.getLogger(__name__)


@dataclasses.dataclass(frozen=True)
class TenantContext:
    """Immutable, serializable tenant identity carried on each message."""

    tenant_id: str
    tenant_name: str
    user_id: str
    user_name: str

    @classmethod
    def from_value(
        cls,
        value: Optional[Union["TenantContext", Dict[str, Any]]],
    ) -> Optional["TenantContext"]:
        """Coerce a tenant-shaped value into ``Optional[TenantContext]``.

        Accepts:
        - ``None`` -> returns ``None``
        - an existing :class:`TenantContext` -> returned as-is
        - a ``dict`` with the four ``tenant_id`` / ``tenant_name`` /
          ``user_id`` / ``user_name`` keys -> reconstructed into a
          ``TenantContext`` (this is the form replay deserialization
          produces before reconstruction has run)
        - a ``dict`` missing required keys, or any other type ->
          returns ``None``

        Callers do ``ctx = TenantContext.from_value(x); ctx.headers()``
        (with a ``None``-check) and don't have to know whether ``x``
        was already a context, a dict, or absent.
        """
        if value is None:
            return None
        if isinstance(value, cls):
            return value
        if isinstance(value, dict):
            try:
                return cls(
                    tenant_id=value["tenant_id"],
                    tenant_name=value["tenant_name"],
                    user_id=value["user_id"],
                    user_name=value["user_name"],
                )
            except KeyError:
                return None
        LOGGER.debug(
            "TenantContext.from_value: unrecognized type %s, treating as None",
            type(value).__name__,
        )
        return None

    def headers(self) -> Dict[str, str]:
        return {
            "X-Tenant-ID": self.tenant_id,
            "X-Tenant-Name": self.tenant_name,
            "X-User-ID": self.user_id,
            "X-User-Name": self.user_name,
        }

    def group_key(self) -> Tuple[str, str]:
        return (self.tenant_id, self.user_id)


TenantResolver = Callable[[], Optional[TenantContext]]


class TenantManager:
    """Process-wide registry for the app's tenant resolver."""

    _resolver: Optional[TenantResolver] = None
    _resolver_error_logged: bool = False

    @classmethod
    def set_resolver(cls, resolver: Optional[TenantResolver]) -> None:
        cls._resolver = resolver
        cls._resolver_error_logged = False

    @classmethod
    def get_current_tenant(cls) -> Optional[TenantContext]:
        if cls._resolver is None:
            return None
        try:
            return cls._resolver()
        except Exception as ex:  # noqa: BLE001 — never break SDK on bad app code
            if not cls._resolver_error_logged:
                LOGGER.warning(
                    "Tenant resolver raised %s: %s. Falling back to no tenant.",
                    type(ex).__name__,
                    ex,
                )
                cls._resolver_error_logged = True
            return None
