# Copyright (c) Metis. All rights reserved.

"""A/B Experiments module for MantisDK.

This module provides SDK-first A/B testing for AI agents. Experiments are
defined directly in Python code, not via a UI, making them easy to version
control and deploy alongside your agent code.

Quick Start::

    import mantisdk as msk

    # Initialize the experiments system with backend connectivity
    msk.init_experiments(
        api_base_url="https://insight.example.com",
        public_key="pk-lf-xxx",
        secret_key="sk-lf-xxx",
    )

    # Register an experiment
    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,
    )

    # Get variant for a user
    variant = msk.experiments.get_variant("prompt-test-v1", user_id="user-alice-789")
    if variant:
        system_prompt = variant.config["prompt"]
        # Run your agent with this prompt
        response = my_agent(query, system_prompt=system_prompt)

How It Works:

1. **Deterministic Assignment**: The same (experiment_key, user_id) pair always
   returns the same variant. This uses a SHA256 hash for uniform distribution.
   For batch experiments where all users are known upfront, use
   ``compute_balanced_assignments`` to guarantee an exact 50/50 split.

2. **Automatic Tracing**: When you call get_variant(), the experiment_id and
   variant_key are automatically injected into all subsequent OTel spans via
   the TracingContextSpanProcessor.

3. **Exposure Logging**: Exposures are logged asynchronously in the background
   and batched for efficiency.

4. **Traffic Control**: Use traffic_percent < 100 to gradually roll out
   experiments to a subset of users.

Exports:
    - Variant: A variant with key, weight, and config
    - ExperimentConfig: An experiment definition
    - Exposure: An exposure record (internal use)
    - ExperimentClient: The main client class
    - get_current_experiment: Get current experiment context (for span injection)
    - clear_current_experiment: Clear current experiment context
"""

from .assignment import assign_variant, compute_balanced_assignments, compute_hash
from .client import (
    ExperimentClient,
    _set_experiments_client,
    clear_current_experiment,
    get_current_experiment,
    get_experiments_client,
)
from .exposure import ExposureLogger
from .types import ExperimentConfig, Exposure, Variant

__all__ = [
    # Types
    "Variant",
    "ExperimentConfig",
    "Exposure",
    # Assignment
    "compute_hash",
    "assign_variant",
    "compute_balanced_assignments",
    # Exposure logging
    "ExposureLogger",
    # Client
    "ExperimentClient",
    "get_current_experiment",
    "clear_current_experiment",
    "get_experiments_client",
]
