"""MLflow tracking integration for the Andlake platform.

Provides helpers to configure the MLflow tracking URI and retrieve a
pre-configured ``MlflowClient``.

Usage::

    from andlake import configure_mlflow
    import mlflow

    configure_mlflow(experiment_name="fraud-detection")
    with mlflow.start_run():
        mlflow.log_metric("auc", 0.95)
"""

from __future__ import annotations

from andlake.config import get_config


def configure_mlflow(
    *,
    experiment_name: str | None = None,
) -> None:
    """Set up MLflow tracking URI and optionally set the active experiment.

    Parameters
    ----------
    experiment_name:
        If provided, calls ``mlflow.set_experiment()`` to create or select
        the named experiment.
    """
    import mlflow

    cfg = get_config()
    mlflow.set_tracking_uri(cfg.mlflow_tracking_uri)

    if experiment_name is not None:
        mlflow.set_experiment(experiment_name)


def get_mlflow_client():
    """Return an ``MlflowClient`` pointed at the platform tracking server.

    Returns
    -------
    mlflow.MlflowClient
    """
    from mlflow import MlflowClient

    cfg = get_config()
    return MlflowClient(tracking_uri=cfg.mlflow_tracking_uri)
