# Copyright (c) Microsoft. All rights reserved.

__version__ = "0.1.6"

from .adapter import *
from .algorithm import *
from .client import DevTaskLoader, MantisdkClient  # deprecated  # type: ignore
from .config import *
from .emitter import *
from .env_var import *
from .execution import *

# Instrumentation (auto-instrumentation for LLM libraries)
from .instrumentation import (
    Instrumentor,
    InstrumentorRegistry,
    MetricSchema,
    instrument,
    instrument_all,
    uninstrument,
    uninstrument_all,
)
from .integrations.insight import (
    InsightIngestionClient,
    InsightIngestionConfig,
)
from .litagent import *

# Run tracking (wandb-style experiment tracking)
from .run import (  # Decorator-based API
    Metric,
    Rollout,
    Run,
    episode,
    experiment,
    finish,
    get_current_rollout,
    get_current_run,
    init_run,
    log,
    run,
    score,
)

# LLMProxy requires litellm[proxy] which has strict dependency pins.
# Make import optional to avoid conflicts for users who don't need this feature.
# Install with: pip install 'mantisdk[proxy]'
try:
    from .llm_proxy import *
except ImportError:
    pass  # LLMProxy not available without litellm[proxy] extras

from .logging import configure_logger  # deprecated  # type: ignore
from .logging import setup as setup_logging  # type: ignore
from .logging import setup_module as setup_module_logging  # type: ignore
from .runner import *
from .server import MantisdkServer  # deprecated  # type: ignore
from .store import *
from .tracer import *

# Trainer requires LLMProxy which requires litellm[proxy]
# Make import optional to avoid conflicts for users who don't need training features.
try:
    from .trainer import *
except ImportError:
    pass  # Trainer not available without litellm[proxy] extras

from .types import *

# =============================================================================
# A/B Experiments — module-level proxy
# =============================================================================


class _ExperimentsNamespace:
    """Proxy that delegates to the global ExperimentClient.

    Allows `msk.experiments.register(...)` and `msk.experiments.get_variant(...)`
    without requiring users to manage an ExperimentClient instance directly.

    If no client has been initialized yet (via `init_experiments()`), calling
    `register()` or `get_variant()` will lazily create a local-only client
    (no exposure logging).
    """

    def _client(self):
        from .experiments.client import ExperimentClient, _set_experiments_client, get_experiments_client

        client = get_experiments_client()
        if client is None:
            client = ExperimentClient()
            _set_experiments_client(client)
        return client

    def register(self, key, variants, traffic_percent=100, sync_to_backend=True):
        """Register an A/B experiment. See ExperimentClient.register()."""
        return self._client().register(key, variants, traffic_percent, sync_to_backend)

    def get_variant(self, experiment_key, user_id):
        """Get variant for a user. See ExperimentClient.get_variant()."""
        return self._client().get_variant(experiment_key, user_id)

    def is_registered(self, experiment_key):
        """Check if an experiment is registered."""
        return self._client().is_registered(experiment_key)

    def list_experiments(self):
        """List all registered experiment keys."""
        return self._client().list_experiments()

    def fetch(self, experiment_key, auto_register=True):
        """Fetch an experiment from the backend. See ExperimentClient.fetch()."""
        return self._client().fetch(experiment_key, auto_register)

    def fetch_all(self, auto_register=True):
        """Fetch all experiments from the backend. See ExperimentClient.fetch_all()."""
        return self._client().fetch_all(auto_register)

    def configure_backend(self, api_base_url, public_key, secret_key):
        """Configure backend connectivity. See ExperimentClient.configure_backend()."""
        return self._client().configure_backend(api_base_url, public_key, secret_key)


# Store proxy under a private name so Python's import system can't clobber it
# when 'mantisdk.experiments' (the subpackage) is later imported internally.
_EXPERIMENTS_PROXY = _ExperimentsNamespace()
experiments = _EXPERIMENTS_PROXY


def init_experiments(
    api_endpoint=None, public_key=None, secret_key=None, batch_size=100, flush_interval=1.0, api_base_url=None
):
    """Initialize the experiments system with backend connectivity.

    Sets up the ExposureLogger (background thread that batches and sends
    exposure events to the backend) and creates the global ExperimentClient.

    Args:
        api_endpoint: Full URL for the exposures batch API
                     (e.g., "https://insight.example.com/api/public/experiments/exposures/batch")
                     If not provided, uses api_base_url to construct it.
        public_key: Insight public key (pk-lf-...)
        secret_key: Insight secret key (sk-lf-...)
        batch_size: Max exposures per batch (default 100)
        flush_interval: Seconds between flushes (default 1.0)
        api_base_url: Base URL for the experiments API (e.g., "https://insight.example.com").
                     Used for fetching/syncing experiments. If not provided, derived from api_endpoint.

    Returns:
        The ExperimentClient instance.
    """
    # The imports above cause Python to register 'mantisdk.experiments' as the
    # subpackage module in sys.modules, which overwrites 'mantisdk.experiments'
    # attribute on this module. Restore the proxy immediately.
    import sys

    from .experiments.client import ExperimentClient, _set_experiments_client
    from .experiments.exposure import ExposureLogger

    _self = sys.modules.get(__name__)
    if _self is not None:
        _self.experiments = _EXPERIMENTS_PROXY

    # Derive api_base_url from api_endpoint if not provided
    if api_base_url is None and api_endpoint:
        # Extract base URL from the full endpoint
        # e.g., "https://insight.example.com/api/public/experiments/exposures/batch"
        # -> "https://insight.example.com"
        from urllib.parse import urlparse

        parsed = urlparse(api_endpoint)
        api_base_url = f"{parsed.scheme}://{parsed.netloc}"

    # Derive api_endpoint from api_base_url if not provided
    if api_endpoint is None and api_base_url:
        api_endpoint = f"{api_base_url}/api/public/experiments/exposures/batch"

    exposure_logger = None
    if api_endpoint and public_key and secret_key:
        exposure_logger = ExposureLogger(
            api_endpoint=api_endpoint,
            public_key=public_key,
            secret_key=secret_key,
            batch_size=batch_size,
            flush_interval=flush_interval,
        )
        exposure_logger.start()

    client = ExperimentClient(
        exposure_logger=exposure_logger,
        api_base_url=api_base_url,
        public_key=public_key,
        secret_key=secret_key,
    )
    _set_experiments_client(client)

    # Transfer any experiments registered before init (use _EXPERIMENTS_PROXY
    # directly — 'experiments' module-level name may be the subpackage at this point)
    existing = _EXPERIMENTS_PROXY._client()
    if existing is not client and hasattr(existing, "_experiments"):
        for key, config in existing._experiments.items():
            client._experiments[key] = config

    return client
