"""Multi-tenant httpx auth for the Acceldata AIO tracer.

INTEGRATES WITH (not a monkey-patch)
-------------------------------------
``opik.hooks.add_httpx_client_hook`` — the official extension point
opik provides to mutate every httpx client it constructs. This module
registers a hook that does ``client.auth = MultiTenantAuth()``. So
unlike ``_streamer_patch`` / ``_batcher_patch`` / ``_processor_patch``
/ ``_replay_patch``, this is a sanctioned integration, not a runtime
patch of opik internals.

Parallel to :mod:`acceldata_aio_tracer.auth`. The single-tenant flow
sets ``client.auth = AcceldataAuth(access_key, secret_key)``; the
multi-tenant flow sets ``client.auth = MultiTenantAuth()`` and routes
each request to its tenant's workspace via the four
``X-Tenant-ID`` / ``X-Tenant-Name`` / ``X-User-ID`` / ``X-User-Name``
headers. The two are mutually exclusive: in multi-tenant mode the
gateway-level ``accessKey``/``secretKey`` check is bypassed (the SDK
talks directly to the Opik backend, which reads tenant identity from
the four X-headers). Calling ``setup_multitenant_auth_hook()`` instead
of ``setup_acceldata_auth_hook()`` is the toggle.

DATA FLOW
---------
1. Request thread (FastAPI handler): app sets a principal in its own
   ContextVar. Opik creates a span message; the streamer patch reads
   the registered :class:`TenantManager` resolver and stamps
   ``message.tenant`` from the principal.
2. Background ``QueueConsumer`` thread: opik dequeues the message and
   calls the (patched) handler in ``OpikMessageProcessor``. The
   processor patch enters ``tenant_request_context(message.tenant)``
   which writes the four headers into the per-thread thread-local
   below.
3. Same background thread, deeper in the call: opik's REST client
   issues an ``httpx.Request`` and httpx invokes
   ``client.auth.auth_flow(request)`` — which is our
   :class:`MultiTenantAuth`. We read the thread-local and merge the
   four headers onto the request before yielding.
4. Same thread, exit of step 2's ``with`` block: the thread-local is
   restored to its previous value (re-entrant-safe).
"""

from __future__ import annotations

import logging
import threading
from contextlib import contextmanager
from typing import Optional

import httpx
import opik.hooks

from .multitenancy import TenantContext

LOGGER = logging.getLogger(__name__)

# Thread-local where the processor patch stashes a tenant's headers for
# the duration of one REST call. Read by ``MultiTenantAuth.auth_flow``
# on the same thread.
_tenant_headers_local = threading.local()

# ---------------------------------------------------------------------
# TODO(remove-when-aio-patched): service-tenant fallback for /is-alive
# health probes
#
# This entire block — _service_tenant, _FALLBACK_URL_PATHS, and the
# corresponding branch inside MultiTenantAuth.auth_flow — exists ONLY
# to work around an opik-backend config bug. The deployed
# `apps/opik-backend/config.yml` `acceldata.exemptPaths` regex was
# `^/is_alive(/.*)?$` (underscore) instead of `^/is-alive(/.*)?$`
# (hyphen). The Python SDK pings `/is-alive/ping`, which therefore
# wasn't exempt from auth, returning 403 on every probe. Opik's
# connection monitor flipped `has_server_connection=False` and stopped
# flushing spans entirely.
#
# The fix is in the backend YAML; once the patched image is deployed
# everywhere this SDK runs against, delete:
# - the _service_tenant module global below
# - the _FALLBACK_URL_PATHS module global below
# - the conditional block inside MultiTenantAuth.auth_flow
# - the service_tenant kwarg on setup_multitenant_auth_hook
# - the matching kwarg threading in __init__.py and _opik_bridge.py
# - the matching kwarg passed from ADM's tracing.py
#
# Optional and opt-in: callers who don't pass service_tenant get the
# vanilla MultiTenantAuth behavior (no fallback, probe will 403 if the
# backend still has the typo).
# ---------------------------------------------------------------------
_service_tenant: Optional[TenantContext] = None

# URL substrings that may use the service-tenant fallback when no
# request-scoped tenant is set. We deliberately keep this list small —
# real trace/span traffic must never silently route to the service
# tenant.
_FALLBACK_URL_PATHS = ("/is-alive/ping",)
# ---------------------------------------------------------------------
# END TODO(remove-when-aio-patched) block 1 of 2 in this file
# ---------------------------------------------------------------------


@contextmanager
def tenant_request_context(tenant: Optional[TenantContext]):
    """Set per-thread tenant headers for the duration of one REST call.

    A ``None`` tenant is a no-op (no thread-local mutation, no headers
    added). Nested usage saves and restores the previous value so the
    helper is re-entrant within a thread.
    """
    if tenant is None:
        yield
        return
    previous = getattr(_tenant_headers_local, "headers", None)
    _tenant_headers_local.headers = tenant.headers()
    try:
        yield
    finally:
        _tenant_headers_local.headers = previous


class MultiTenantAuth(httpx.Auth):
    """httpx.Auth implementation that injects the four tenant headers
    from a per-thread :class:`threading.local`.

    Replaces :class:`acceldata_aio_tracer.auth.AcceldataAuth` in
    multi-tenant deployments. When no tenant headers are set on the
    current thread (e.g. for non-trace/span requests, or for traffic
    that originated outside a ``tenant_request_context`` block), no
    headers are injected — the request goes out as the SDK constructed
    it.
    """

    def auth_flow(self, request):  # type: ignore[override]
        tenant_hdrs = getattr(_tenant_headers_local, "headers", None)

        # ---- TODO(remove-when-aio-patched) block 2 of 2 in this file ----
        # See the comment block above _service_tenant for full context.
        # Service-tenant fallback for /is-alive/ping probes; remove
        # along with the module globals once the backend YAML is fixed.
        if tenant_hdrs is None and _service_tenant is not None:
            url_path = str(request.url.path) if hasattr(request.url, "path") else str(request.url)
            if any(p in url_path for p in _FALLBACK_URL_PATHS):
                tenant_hdrs = _service_tenant.headers()
        # ---- END TODO(remove-when-aio-patched) block 2 of 2 ----

        if tenant_hdrs:
            for k, v in tenant_hdrs.items():
                # Only set if the caller didn't already supply one — preserves
                # the explicit-headers-win semantic the test suite relies on.
                if k not in request.headers:
                    request.headers[k] = v
        yield request


def setup_multitenant_auth_hook(
    service_tenant: Optional[TenantContext] = None,
) -> None:
    """Register an opik httpx-client hook that installs
    :class:`MultiTenantAuth` on every Opik httpx client.

    Mirrors :func:`acceldata_aio_tracer.auth.setup_acceldata_auth_hook`
    but for multi-tenant routing instead of single-tenant apiKey auth.
    Idempotent — safe to call multiple times. The hook is process-wide;
    every Opik client instantiated after this call carries
    ``MultiTenantAuth``.

    Args:
        service_tenant: Optional fallback :class:`TenantContext` for
            requests that no per-request ``tenant_request_context``
            covers — currently only the ``/is-alive/ping`` health
            probe. See the ``_service_tenant`` block at the top of
            this file for full context.
    """
    global _service_tenant
    if service_tenant is not None:
        _service_tenant = service_tenant

    def _multitenant_client_hook(client: httpx.Client) -> None:
        client.auth = MultiTenantAuth()

    opik.hooks.add_httpx_client_hook(
        opik.hooks.HttpxClientHook(
            client_modifier=_multitenant_client_hook,
            client_init_arguments=None,
        )
    )
