# Copyright (c) Metis. All rights reserved.

"""Main client API for A/B experiments.

This module provides the ExperimentClient class that users interact with
to register experiments and get variant assignments. It also manages the
context variable that the span processor reads from.

The client supports two modes:
1. **Local-only**: Experiments registered in code stay in-memory
2. **Backend-synced**: Experiments are fetched from and synced to the Insight backend

Usage::

    import mantisdk as msk

    # Initialize with backend sync (recommended)
    msk.init_experiments(
        api_endpoint="https://insight.example.com/api/public/experiments/exposures/batch",
        public_key="pk-lf-...",
        secret_key="sk-lf-...",
    )

    # Option A: Register experiment in code (syncs to backend in background)
    msk.experiments.register(
        key="prompt-test-v1",
        variants=[
            {"key": "control", "weight": 50, "config": {"prompt": "..."}},
            {"key": "treatment", "weight": 50, "config": {"prompt": "..."}},
        ],
        traffic_percent=100,
    )

    # Option B: Fetch experiment from backend (defined in UI)
    msk.experiments.fetch("prompt-test-v1")

    # Get variant for a user (~0ms, purely in-memory)
    variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-123")
    if variant:
        print(f"Assigned to {variant.key}")
        system_prompt = variant.config["prompt"]
"""

from __future__ import annotations

import logging
import threading
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional

import requests

from .assignment import assign_variant
from .types import ExperimentConfig, Exposure, Variant

if TYPE_CHECKING:
    from .exposure import ExposureLogger

logger = logging.getLogger(__name__)


# =============================================================================
# Exceptions
# =============================================================================


class BalancedAssignmentError(Exception):
    """Raised when server-side balanced assignment fails due to network/server issues.

    This exception is raised by get_variant_balanced() when the server cannot be reached
    or returns an error. The caller (get_variant) catches this and falls back to
    local hash-based assignment for graceful degradation.

    This is NOT raised when the user is legitimately not in traffic (server returns
    variantKey: null) — in that case, None is returned normally.
    """

    pass


# =============================================================================
# Context Variable for Experiment Assignment
# =============================================================================

# Stores the current experiment context for span injection.
# The TracingContextSpanProcessor reads this and injects attributes into spans.
_current_experiment: ContextVar[Optional[Dict[str, Any]]] = ContextVar("msk_current_experiment", default=None)


def get_current_experiment() -> Optional[Dict[str, Any]]:
    """Get the current experiment assignment context.

    Returns a dict with keys:
    - experiment_id: The experiment key
    - variant_key: The assigned variant key
    - user_id: The user identifier

    Returns None if no experiment is active.

    This is called by TracingContextSpanProcessor to inject experiment
    attributes into spans.
    """
    return _current_experiment.get()


def _set_current_experiment(ctx: Optional[Dict[str, Any]]) -> None:
    """Set the current experiment context.

    Internal function - called by ExperimentClient.get_variant().
    """
    _current_experiment.set(ctx)


def clear_current_experiment() -> None:
    """Clear the current experiment context.

    Call this when you want to stop attributing traces to an experiment,
    for example when switching users or ending an experiment session.
    """
    _current_experiment.set(None)


# =============================================================================
# Experiment Client
# =============================================================================


class ExperimentClient:
    """Client for managing A/B experiments.

    This class provides the main API for:
    - Registering experiments with variants (locally and to backend)
    - Fetching experiments from the backend
    - Getting variant assignments for users (~0ms, in-memory)
    - Managing exposure logging (async, background thread)
    - Background config polling (keeps in-memory cache fresh)

    The client is typically accessed via the mantisdk module:
        msk.experiments.register(...)
        msk.experiments.fetch(...)
        msk.experiments.get_variant(...)

    Attributes:
        exposure_logger: The background logger for exposure events
        api_base_url: Base URL for the experiments API
        public_key: Insight public key for authentication
        secret_key: Insight secret key for authentication
        poll_interval: Seconds between background config refreshes (default 30)
    """

    def __init__(
        self,
        exposure_logger: Optional["ExposureLogger"] = None,
        api_base_url: Optional[str] = None,
        public_key: Optional[str] = None,
        secret_key: Optional[str] = None,
        poll_interval: float = 30.0,
    ):
        """Initialize the experiment client.

        Args:
            exposure_logger: Optional ExposureLogger for sending exposure events.
                            If not provided, exposures will be logged but not sent.
            api_base_url: Base URL for the experiments API (e.g., "https://insight.example.com").
                         If not provided, backend sync is disabled.
            public_key: Insight public key for API authentication.
            secret_key: Insight secret key for API authentication.
            poll_interval: Seconds between background polls for config updates.
                          Default is 30 seconds. Set to 0 to disable polling.
        """
        # PREVIOUSLY: plain dict, written from main thread only.
        # NOW: protected by a lock because the background refresh thread also
        # writes to it while get_variant() reads from the main thread.
        self._experiments: Dict[str, ExperimentConfig] = {}
        self._experiments_lock = threading.Lock()

        self._exposure_logger = exposure_logger
        self._api_base_url = api_base_url
        self._public_key = public_key
        self._secret_key = secret_key
        self._poll_interval = poll_interval

        # Background polling state
        self._poll_thread: Optional[threading.Thread] = None
        self._poll_shutdown = threading.Event()
        self._poll_started = False

    def set_exposure_logger(self, logger: "ExposureLogger") -> None:
        """Set the exposure logger after initialization.

        This is useful when the client is created before the API configuration
        is known (e.g., during module initialization).
        """
        self._exposure_logger = logger

    def configure_backend(
        self,
        api_base_url: str,
        public_key: str,
        secret_key: str,
    ) -> None:
        """Configure backend connectivity for fetching/syncing experiments.

        Args:
            api_base_url: Base URL for the experiments API.
            public_key: Insight public key.
            secret_key: Insight secret key.
        """
        self._api_base_url = api_base_url
        self._public_key = public_key
        self._secret_key = secret_key
        logger.debug(f"Configured experiments backend: {api_base_url}")

    def _get_auth_header(self) -> Dict[str, str]:
        """Get authorization header for API requests."""
        if not self._public_key or not self._secret_key:
            return {}
        import base64

        credentials = base64.b64encode(f"{self._public_key}:{self._secret_key}".encode()).decode()
        return {"Authorization": f"Basic {credentials}"}

    @staticmethod
    def _parse_experiment_response(data: Dict[str, Any]) -> ExperimentConfig:
        """Parse an experiment JSON response into an ExperimentConfig."""
        return ExperimentConfig(
            key=data["key"],
            variants=[
                Variant(
                    key=v["key"],
                    weight=v["weight"],
                    config=v.get("config", {}),
                )
                for v in data["variants"]
            ],
            traffic_percent=data["trafficPercent"],
        )

    # -------------------------------------------------------------------------
    # Background Config Polling
    # -------------------------------------------------------------------------

    def start_background_refresh(self) -> None:
        """Start a background daemon thread that polls the backend for config updates.

        WHY: Previously, config was only loaded when the caller explicitly called
        fetch() or fetch_all() — a blocking HTTP call. Now, a daemon thread runs
        silently in the background and atomically swaps in fresh configs every
        poll_interval seconds. get_variant() always reads from warm memory (~0ms).

        Safe to call multiple times — only starts once. No-op if poll_interval=0
        or backend is not configured.
        """
        if self._poll_started or not self._api_base_url or self._poll_interval <= 0:
            return

        self._poll_shutdown.clear()
        self._poll_thread = threading.Thread(
            target=self._poll_loop,
            daemon=True,
            name="msk-experiment-refresh",
        )
        self._poll_thread.start()
        self._poll_started = True
        logger.debug(f"Background experiment refresh started (interval={self._poll_interval}s)")

    def stop_background_refresh(self) -> None:
        """Stop the background refresh thread. Blocks until it exits (up to 5s)."""
        if not self._poll_started:
            return
        self._poll_shutdown.set()
        if self._poll_thread:
            self._poll_thread.join(timeout=5.0)
        self._poll_started = False
        logger.debug("Background experiment refresh stopped")

    def _poll_loop(self) -> None:
        """Background loop: refresh all known experiment keys every poll_interval."""
        # Run immediately on first iteration, then wait between subsequent ones.
        while not self._poll_shutdown.is_set():
            self._refresh_all_from_backend()
            # wait() returns True if shutdown was requested, False on timeout.
            # Using wait() instead of sleep() so shutdown is instant.
            self._poll_shutdown.wait(timeout=self._poll_interval)

    def _refresh_all_from_backend(self) -> None:
        """Silently re-fetch all registered experiment keys from the backend.

        Called by the background poll loop. Updates _experiments atomically
        under a lock so get_variant() never sees a partial state.
        """
        if not self._api_base_url:
            return

        with self._experiments_lock:
            keys = list(self._experiments.keys())

        if not keys:
            return

        headers = self._get_auth_header()
        updated = 0

        for key in keys:
            url = f"{self._api_base_url}/api/public/experiments/{key}"
            try:
                response = requests.get(url, headers=headers, timeout=5)
                if response.status_code == 200:
                    config = self._parse_experiment_response(response.json())
                    with self._experiments_lock:
                        self._experiments[key] = config
                    updated += 1
                elif response.status_code == 404:
                    logger.debug(f"Experiment '{key}' no longer found in backend")
                else:
                    logger.debug(f"Background refresh for '{key}' got {response.status_code}")
            except Exception as e:
                # Never crash the background thread — just log and move on.
                logger.debug(f"Background refresh failed for '{key}': {e}")

        if updated:
            logger.debug(f"Background refresh updated {updated}/{len(keys)} experiments")

    # -------------------------------------------------------------------------
    # Register / Sync
    # -------------------------------------------------------------------------

    def register(
        self,
        key: str,
        variants: List[Dict[str, Any]],
        traffic_percent: int = 100,
        sync_to_backend: bool = True,
    ) -> None:
        """Register an experiment configuration.

        Experiments are defined in code and can optionally be synced to the
        backend database. This makes them visible in the UI and allows for
        evaluator configuration.

        PREVIOUSLY: if sync_to_backend=True, this blocked the caller for the
        full HTTP round-trip (~100–500ms) before returning.
        NOW: the backend sync is fire-and-forget (daemon thread). register()
        returns immediately after writing to in-memory state.

        Args:
            key: Unique experiment identifier (e.g., "prompt-test-v1").
                 Should be descriptive and versioned.
            variants: List of variant dicts, each with:
                - key: Variant identifier (e.g., "control", "treatment")
                - weight: Relative weight for traffic (e.g., 50 = 50%)
                - config: Arbitrary config dict passed to your code
            traffic_percent: Percentage of users to include (0-100).
                            Use < 100 to gradually roll out experiments.
            sync_to_backend: If True and backend is configured, syncs the
                            experiment to the Insight database. Default True.

        Example::

            msk.experiments.register(
                key="prompt-test-v1",
                variants=[
                    {"key": "control", "weight": 50, "config": {"prompt": "You are a helpful assistant."}},
                    {"key": "treatment", "weight": 50, "config": {"prompt": "You are a concise assistant."}},
                ],
                traffic_percent=100,
            )
        """
        if traffic_percent < 0 or traffic_percent > 100:
            raise ValueError(f"traffic_percent must be 0-100, got {traffic_percent}")

        if len(variants) < 2:
            raise ValueError("Experiments require at least 2 variants")

        config = ExperimentConfig(
            key=key,
            variants=[
                Variant(
                    key=v["key"],
                    weight=v.get("weight", 1),
                    config=v.get("config", {}),
                )
                for v in variants
            ],
            traffic_percent=traffic_percent,
        )

        with self._experiments_lock:
            self._experiments[key] = config

        logger.debug(f"Registered experiment '{key}' with {len(variants)} variants, " f"{traffic_percent}% traffic")

        # PREVIOUSLY: blocking requests.post() here — caller waited for it.
        # NOW: spawns a daemon thread and returns immediately.
        if sync_to_backend and self._api_base_url:
            self._sync_experiment_to_backend_async(key, variants, traffic_percent)

    def _sync_experiment_to_backend_async(
        self,
        key: str,
        variants: List[Dict[str, Any]],
        traffic_percent: int,
    ) -> None:
        """Fire-and-forget backend sync — runs in a daemon thread.

        PREVIOUSLY: this was _sync_experiment_to_backend() and was called
        directly (blocking). Now it spawns a thread and returns in ~0ms.
        """

        def _do_sync() -> None:
            url = f"{self._api_base_url}/api/public/experiments"
            headers = {
                "Content-Type": "application/json",
                **self._get_auth_header(),
            }
            payload = {
                "key": key,
                "trafficPercent": traffic_percent,
                "variants": [
                    {
                        "key": v["key"],
                        "weight": v.get("weight", 1),
                        "config": v.get("config", {}),
                    }
                    for v in variants
                ],
            }
            try:
                response = requests.post(url, json=payload, headers=headers, timeout=5)
                if response.status_code in (200, 201):
                    data = response.json()
                    action = "created" if data.get("created") else "updated"
                    logger.info(f"Experiment '{key}' {action} in backend (id={data.get('id')})")
                else:
                    logger.warning(
                        f"Failed to sync experiment '{key}' to backend: " f"{response.status_code} {response.text}"
                    )
            except Exception as e:
                logger.warning(f"Failed to sync experiment '{key}' to backend: {e}")

        thread = threading.Thread(
            target=_do_sync,
            daemon=True,
            name=f"msk-sync-{key}",
        )
        thread.start()

    # -------------------------------------------------------------------------
    # Fetch (explicit blocking fetch — for startup patterns)
    # -------------------------------------------------------------------------

    def fetch(
        self,
        experiment_key: str,
        auto_register: bool = True,
    ) -> Optional[ExperimentConfig]:
        """Fetch an experiment configuration from the backend.

        This is intentionally synchronous — use it at startup when you need
        the config to be available before proceeding. For runtime updates,
        rely on start_background_refresh() instead.

        Args:
            experiment_key: The experiment identifier to fetch.
            auto_register: If True, automatically register the fetched
                          experiment locally so get_variant() works.

        Returns:
            The ExperimentConfig if found, None otherwise.

        Example::

            # Fetch experiment defined in the UI
            config = msk.experiments.fetch("prompt-test-v1")
            if config:
                variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-123")
        """
        if not self._api_base_url:
            logger.warning("Backend not configured, cannot fetch experiment")
            return None

        url = f"{self._api_base_url}/api/public/experiments/{experiment_key}"
        headers = self._get_auth_header()

        try:
            response = requests.get(url, headers=headers, timeout=5)
            if response.status_code == 200:
                config = self._parse_experiment_response(response.json())

                if auto_register:
                    with self._experiments_lock:
                        self._experiments[experiment_key] = config
                    logger.info(
                        f"Fetched and registered experiment '{experiment_key}' "
                        f"from backend ({len(config.variants)} variants)"
                    )
                else:
                    logger.info(f"Fetched experiment '{experiment_key}' from backend")

                return config
            elif response.status_code == 404:
                logger.warning(f"Experiment '{experiment_key}' not found in backend")
                return None
            else:
                logger.warning(
                    f"Failed to fetch experiment '{experiment_key}': " f"{response.status_code} {response.text}"
                )
                return None
        except Exception as e:
            logger.warning(f"Failed to fetch experiment '{experiment_key}': {e}")
            return None

    def fetch_all(self, auto_register: bool = True) -> List[ExperimentConfig]:
        """Fetch all active experiments from the backend.

        The listing endpoint returns summaries only (no variant details), so
        each experiment must be fetched individually via ``fetch()``.

        Useful at startup to pre-populate the in-memory cache. After calling
        this, ``start_background_refresh()`` will keep configs fresh.

        Args:
            auto_register: If True, automatically register fetched experiments.

        Returns:
            List of ExperimentConfig objects that were successfully fetched.
        """
        if not self._api_base_url:
            logger.warning("Backend not configured, cannot fetch experiments")
            return []

        url = f"{self._api_base_url}/api/public/experiments"
        headers = self._get_auth_header()

        try:
            response = requests.get(url, headers=headers, timeout=10)
            if response.status_code != 200:
                logger.warning(f"Failed to fetch experiments: {response.status_code} {response.text}")
                return []

            experiment_keys = [exp["key"] for exp in response.json().get("experiments", [])]
        except Exception as e:
            logger.warning(f"Failed to fetch experiments: {e}")
            return []

        configs: List[ExperimentConfig] = []
        for key in experiment_keys:
            config = self.fetch(key, auto_register=auto_register)
            if config:
                configs.append(config)

        logger.info(f"Fetched {len(configs)}/{len(experiment_keys)} experiments from backend")
        return configs

    # -------------------------------------------------------------------------
    # Get Variant (~0ms hot path)
    # -------------------------------------------------------------------------

    def get_variant(
        self,
        experiment_key: str,
        user_id: str,
    ) -> Optional[Variant]:
        """Get the variant for a user in an experiment.

        This is the main entry point for A/B testing. It:
        1. Deterministically assigns the user to a variant (or None if not in traffic)
        2. Sets the context variable for span injection
        3. Logs an exposure event (sent to backend asynchronously)

        When the backend is configured, uses server-side balanced assignment
        (Redis counter) for exact proportional allocation at any sample size.
        Falls back to local hash-based assignment when the backend is unavailable.

        Args:
            experiment_key: The experiment identifier (must be registered first)
            user_id: The user/subject identifier. This should be a stable ID
                    that uniquely identifies the user across sessions.

        Returns:
            The assigned Variant with its config, or None if:
            - The experiment is not registered
            - The user is not in the traffic allocation

        Example::

            variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-alice-789")
            if variant:
                print(f"User assigned to {variant.key}")
                system_prompt = variant.config["prompt"]
                # Run your agent with this prompt
            else:
                # User not in experiment, use default behavior
                system_prompt = DEFAULT_PROMPT
        """
        if self._api_base_url:
            try:
                result = self.get_variant_balanced(experiment_key, user_id)
                # result is None only when user is legitimately not in traffic
                # (server returned variantKey: null) — do NOT fall back
                return result
            except BalancedAssignmentError:
                # Network error, timeout, server error — fall back to hash-based
                logger.warning(
                    f"Balanced assignment failed for '{experiment_key}', " f"falling back to hash-based assignment"
                )

        return self._get_variant_hash(experiment_key, user_id)

    def _get_variant_hash(
        self,
        experiment_key: str,
        user_id: str,
    ) -> Optional[Variant]:
        """Hash-based local assignment (fallback when backend is unavailable)."""
        with self._experiments_lock:
            config = self._experiments.get(experiment_key)

        if config is None:
            logger.warning(
                f"Experiment '{experiment_key}' not registered. "
                f"Call msk.experiments.register() or msk.experiments.fetch() first."
            )
            return None

        variant = assign_variant(config, user_id)
        if variant is None:
            logger.debug(f"User '{user_id}' not in traffic for '{experiment_key}'")
            return None

        _set_current_experiment(
            {
                "experiment_id": experiment_key,
                "variant_key": variant.key,
                "user_id": user_id,
            }
        )

        if self._exposure_logger:
            self._exposure_logger.log(
                Exposure(
                    user_id=user_id,
                    experiment_id=experiment_key,
                    variant_key=variant.key,
                    timestamp=datetime.now(timezone.utc).isoformat(),
                )
            )
        else:
            logger.debug(
                f"No exposure logger configured, exposure not sent "
                f"(experiment={experiment_key}, variant={variant.key}, user={user_id})"
            )

        logger.debug(f"Assigned user '{user_id}' to variant '{variant.key}' " f"in experiment '{experiment_key}'")

        return variant

    def assign_with_override(
        self,
        experiment_key: str,
        user_id: str,
        variant_key: str,
    ) -> Optional[Variant]:
        """Assign a user to a specific variant, bypassing hash-based selection.

        Use this when you have pre-computed balanced assignments (e.g. via
        ``compute_balanced_assignments``) and want to force a specific variant
        while still setting the tracing context and logging the exposure.

        Args:
            experiment_key: The experiment identifier (must be registered first).
            user_id: The user/subject identifier.
            variant_key: The variant key to force (e.g. "control", "treatment").

        Returns:
            The matching Variant with its config, or None if the experiment is
            not registered or variant_key is not found.

        Example::

            assignments = compute_balanced_assignments(config, all_users)
            for user_id, variant in assignments.items():
                if variant:
                    msk.experiments.assign_with_override(
                        "my-exp", user_id, variant.key
                    )
        """
        with self._experiments_lock:
            config = self._experiments.get(experiment_key)

        if config is None:
            logger.warning(
                f"Experiment '{experiment_key}' not registered. "
                f"Call msk.experiments.register() or msk.experiments.fetch() first."
            )
            return None

        variant = next((v for v in config.variants if v.key == variant_key), None)
        if variant is None:
            logger.warning(f"Variant '{variant_key}' not found in experiment '{experiment_key}'")
            return None

        _set_current_experiment(
            {
                "experiment_id": experiment_key,
                "variant_key": variant.key,
                "user_id": user_id,
            }
        )

        if self._exposure_logger:
            self._exposure_logger.log(
                Exposure(
                    user_id=user_id,
                    experiment_id=experiment_key,
                    variant_key=variant.key,
                    timestamp=datetime.now(timezone.utc).isoformat(),
                )
            )

        logger.debug(f"Force-assigned user '{user_id}' to variant '{variant.key}' " f"in experiment '{experiment_key}'")

        return variant

    def get_variant_balanced(
        self,
        experiment_key: str,
        user_id: str,
    ) -> Optional[Variant]:
        """Get variant using server-side balanced assignment (guaranteed 50/50).

        Unlike ``get_variant`` which uses local hash-based assignment (can produce
        unequal splits at small sample sizes), this method calls the server which
        uses a Redis atomic counter to guarantee exact proportional allocation.

        Trade-off: requires a network call (~10-50ms) vs hash-based (~0ms local).
        Use this when you have small sample sizes and need exact balance.

        The assignment is:
        1. Deterministic per user — same user always gets same variant
        2. Perfectly balanced — guaranteed 50/50 (or proportional to weights)
        3. Persistent — stored in Redis, survives restarts

        Args:
            experiment_key: The experiment identifier (must exist in backend).
            user_id: The user/subject identifier.

        Returns:
            The assigned Variant with its config, or None if:
            - User is legitimately not in traffic (server returned variantKey: null)

        Raises:
            BalancedAssignmentError: If network error, timeout, or server error occurs.
                Caller can catch this to fall back to hash-based assignment.

        Example::

            # Use balanced assignment for small-sample experiments
            variant = msk.experiments.get_variant_balanced(
                "prompt-test-v1", user_id="user-alice-789"
            )
            if variant:
                print(f"User assigned to {variant.key} (server-balanced)")
        """
        if not self._api_base_url:
            logger.warning(
                "Backend not configured — cannot use balanced assignment. "
                "Falling back to local hash-based assignment."
            )
            return self.get_variant(experiment_key, user_id)

        url = f"{self._api_base_url}/api/public/experiments/assign"
        headers = self._get_auth_header()
        payload = {
            "experimentKey": experiment_key,
            "userId": user_id,
        }

        try:
            response = requests.post(url, json=payload, headers=headers, timeout=10)

            if response.status_code != 200:
                error_msg = response.text[:200] if response.text else "Unknown error"
                logger.error(f"Server-balanced assignment failed: {response.status_code} — {error_msg}")
                raise BalancedAssignmentError(f"Server returned {response.status_code}: {error_msg}")

            data = response.json()

            if data.get("variantKey") is None:
                logger.debug(f"User '{user_id}' not in traffic for '{experiment_key}' (server response)")
                return None

            variant = Variant(
                key=data["variantKey"],
                weight=0,  # Weight not returned by assignment endpoint
                config=data.get("config", {}),
            )

            # Set context for span processor
            _set_current_experiment(
                {
                    "experiment_id": experiment_key,
                    "variant_key": variant.key,
                    "user_id": user_id,
                }
            )

            # Log exposure
            if self._exposure_logger:
                self._exposure_logger.log(
                    Exposure(
                        user_id=user_id,
                        experiment_id=experiment_key,
                        variant_key=variant.key,
                        timestamp=datetime.now(timezone.utc).isoformat(),
                    )
                )

            is_new = data.get("isNewAssignment", True)
            logger.debug(
                f"Server-balanced assignment: user '{user_id}' → '{variant.key}' "
                f"in experiment '{experiment_key}' (new={is_new})"
            )

            return variant

        except requests.exceptions.Timeout:
            logger.error(f"Server-balanced assignment timed out for '{experiment_key}'")
            raise BalancedAssignmentError(f"Timeout for '{experiment_key}'")
        except requests.exceptions.RequestException as e:
            logger.error(f"Server-balanced assignment failed for '{experiment_key}': {e}")
            raise BalancedAssignmentError(f"Network error for '{experiment_key}': {e}")
        except BalancedAssignmentError:
            raise  # Re-raise our custom exception
        except Exception as e:
            logger.error(f"Unexpected error in server-balanced assignment: {e}")
            raise BalancedAssignmentError(f"Unexpected error for '{experiment_key}': {e}")

    # -------------------------------------------------------------------------
    # Introspection
    # -------------------------------------------------------------------------

    def is_registered(self, experiment_key: str) -> bool:
        """Check if an experiment is registered."""
        with self._experiments_lock:
            return experiment_key in self._experiments

    def list_experiments(self) -> List[str]:
        """List all registered experiment keys."""
        with self._experiments_lock:
            return list(self._experiments.keys())


# =============================================================================
# Module-level instance (populated by mantisdk.init())
# =============================================================================

# This will be set by the main mantisdk module during init()
_experiments_client: Optional[ExperimentClient] = None


def get_experiments_client() -> Optional[ExperimentClient]:
    """Get the global experiments client instance.

    Returns None if mantisdk has not been initialized.
    """
    return _experiments_client


def _set_experiments_client(client: Optional[ExperimentClient]) -> None:
    """Set the global experiments client instance.

    Internal function - called by mantisdk.init().
    """
    global _experiments_client
    _experiments_client = client
