"""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

import os

from andlake.config import get_config


def _resolve_token(access_token: str | None) -> str | None:
    """Resolve an access token from (in priority order):

    1. Explicit ``access_token`` argument
    2. ``ANDLAKE_ACCESS_TOKEN`` environment variable
    3. Keycloak client-credentials grant (auto-fetch)
    """
    if access_token:
        return access_token

    cfg = get_config()
    if cfg.access_token:
        return cfg.access_token

    # Auto-fetch from Keycloak if client credentials are configured
    if cfg.keycloak_client_id and cfg.keycloak_client_secret:
        from andlake.auth import get_access_token
        return get_access_token()

    return None


def _set_tracking_token(access_token: str | None) -> None:
    """Inject a Bearer token for MLflow's REST client.

    MLflow reads ``MLFLOW_TRACKING_TOKEN`` and sends it as an
    ``Authorization: Bearer <token>`` header on every request.
    """
    token = _resolve_token(access_token)
    if token:
        os.environ["MLFLOW_TRACKING_TOKEN"] = token


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

    Authentication is resolved automatically:

    1. If ``access_token`` is passed, use it.
    2. Else if ``ANDLAKE_ACCESS_TOKEN`` env var is set, use it.
    3. Else if Keycloak client credentials are configured
       (``ANDLAKE_KEYCLOAK_CLIENT_ID`` / ``ANDLAKE_KEYCLOAK_CLIENT_SECRET``),
       fetch a token via client-credentials grant.

    Parameters
    ----------
    experiment_name:
        If provided, calls ``mlflow.set_experiment()`` to create or select
        the named experiment.
    access_token:
        Keycloak JWT to authenticate with MLflow.
    """
    import mlflow

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

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


def get_mlflow_client(*, access_token: str | None = None):
    """Return an ``MlflowClient`` pointed at the platform tracking server.

    Parameters
    ----------
    access_token:
        Keycloak JWT to authenticate with MLflow.

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

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