"""
The ``mlflow`` module provides a high-level "fluent" API for starting and managing MLflow runs.
For example:

.. code:: python

    import mlflow

    mlflow.start_run()
    mlflow.log_param("my", "param")
    mlflow.log_metric("score", 100)
    mlflow.end_run()

You can also use the context manager syntax like this:

.. code:: python

    with mlflow.start_run() as run:
        mlflow.log_param("my", "param")
        mlflow.log_metric("score", 100)

which automatically terminates the run at the end of the ``with`` block.

The fluent tracking API is not currently threadsafe. Any concurrent callers to the tracking API must
implement mutual exclusion manually.

For a lower level API, see the :py:mod:`mlflow.client` module.
"""

import contextlib
from typing import TYPE_CHECKING

from mlflow.version import IS_TRACING_SDK_ONLY, VERSION

__version__ = VERSION

import mlflow.mismatch

# `check_version_mismatch` must be called here before importing any other modules
with contextlib.suppress(Exception):
    mlflow.mismatch._check_version_mismatch()

if not IS_TRACING_SDK_ONLY:
    from mlflow import (
        artifacts,  # noqa: F401
        client,  # noqa: F401
        config,  # noqa: F401
        data,  # noqa: F401
        exceptions,  # noqa: F401
        genai,  # noqa: F401
        models,  # noqa: F401
        projects,  # noqa: F401
        tracking,  # noqa: F401
    )

from mlflow import tracing  # noqa: F401
from mlflow.environment_variables import MLFLOW_CONFIGURE_LOGGING
from mlflow.exceptions import MlflowException
from mlflow.utils.lazy_load import LazyLoader
from mlflow.utils.logging_utils import (
    _configure_mlflow_loggers,
    _install_sensitive_query_param_filter,
)

# Lazily load mlflow flavors to avoid excessive dependencies.
anthropic = LazyLoader("mlflow.anthropic", globals(), "mlflow.anthropic")
ag2 = LazyLoader("mlflow.ag2", globals(), "mlflow.ag2")
agno = LazyLoader("mlflow.agno", globals(), "mlflow.agno")
autogen = LazyLoader("mlflow.autogen", globals(), "mlflow.autogen")
bedrock = LazyLoader("mlflow.bedrock", globals(), "mlflow.bedrock")
catboost = LazyLoader("mlflow.catboost", globals(), "mlflow.catboost")
crewai = LazyLoader("mlflow.crewai", globals(), "mlflow.crewai")
diffusers = LazyLoader("mlflow.diffusers", globals(), "mlflow.diffusers")
dspy = LazyLoader("mlflow.dspy", globals(), "mlflow.dspy")
gemini = LazyLoader("mlflow.gemini", globals(), "mlflow.gemini")
groq = LazyLoader("mlflow.groq", globals(), "mlflow.groq")
h2o = LazyLoader("mlflow.h2o", globals(), "mlflow.h2o")
haystack = LazyLoader("mlflow.haystack", globals(), "mlflow.haystack")
johnsnowlabs = LazyLoader("mlflow.johnsnowlabs", globals(), "mlflow.johnsnowlabs")
keras = LazyLoader("mlflow.keras", globals(), "mlflow.keras")
langchain = LazyLoader("mlflow.langchain", globals(), "mlflow.langchain")
lightgbm = LazyLoader("mlflow.lightgbm", globals(), "mlflow.lightgbm")
litellm = LazyLoader("mlflow.litellm", globals(), "mlflow.litellm")
llama_index = LazyLoader("mlflow.llama_index", globals(), "mlflow.llama_index")
metrics = LazyLoader("mlflow.metrics", globals(), "mlflow.metrics")
mistral = LazyLoader("mlflow.mistral", globals(), "mlflow.mistral")
onnx = LazyLoader("mlflow.onnx", globals(), "mlflow.onnx")
otel = LazyLoader("mlflow.otel", globals(), "mlflow.otel")
openai = LazyLoader("mlflow.openai", globals(), "mlflow.openai")
paddle = LazyLoader("mlflow.paddle", globals(), "mlflow.paddle")
pmdarima = LazyLoader("mlflow.pmdarima", globals(), "mlflow.pmdarima")
prophet = LazyLoader("mlflow.prophet", globals(), "mlflow.prophet")
pydantic_ai = LazyLoader("mlflow.pydantic_ai", globals(), "mlflow.pydantic_ai")
pyfunc = LazyLoader("mlflow.pyfunc", globals(), "mlflow.pyfunc")
pyspark = LazyLoader("mlflow.pyspark", globals(), "mlflow.pyspark")
pytorch = LazyLoader("mlflow.pytorch", globals(), "mlflow.pytorch")
rfunc = LazyLoader("mlflow.rfunc", globals(), "mlflow.rfunc")
semantic_kernel = LazyLoader("mlflow.semantic_kernel", globals(), "mlflow.semantic_kernel")
sentence_transformers = LazyLoader(
    "mlflow.sentence_transformers",
    globals(),
    "mlflow.sentence_transformers",
)
shap = LazyLoader("mlflow.shap", globals(), "mlflow.shap")
sklearn = LazyLoader("mlflow.sklearn", globals(), "mlflow.sklearn")
smolagents = LazyLoader("mlflow.smolagents", globals(), "mlflow.smolagents")
spacy = LazyLoader("mlflow.spacy", globals(), "mlflow.spacy")
strands = LazyLoader("mlflow.strands", globals(), "mlflow.strands")
spark = LazyLoader("mlflow.spark", globals(), "mlflow.spark")
statsmodels = LazyLoader("mlflow.statsmodels", globals(), "mlflow.statsmodels")
tensorflow = LazyLoader("mlflow.tensorflow", globals(), "mlflow.tensorflow")
# TxtAI integration is defined at https://github.com/neuml/mlflow-txtai
txtai = LazyLoader("mlflow.txtai", globals(), "mlflow_txtai")
transformers = LazyLoader("mlflow.transformers", globals(), "mlflow.transformers")
xgboost = LazyLoader("mlflow.xgboost", globals(), "mlflow.xgboost")

if TYPE_CHECKING:
    # Do not move this block above the lazy-loaded modules above.
    # All the lazy-loaded modules above must be imported here for code completion to work in IDEs.
    from mlflow import (  # noqa: F401
        ag2,
        agno,
        anthropic,
        autogen,
        bedrock,
        catboost,
        crewai,
        diffusers,
        dspy,
        gemini,
        groq,
        h2o,
        haystack,
        johnsnowlabs,
        keras,
        langchain,
        lightgbm,
        litellm,
        llama_index,
        metrics,
        mistral,
        onnx,
        openai,
        otel,
        paddle,
        pmdarima,
        prophet,
        pydantic_ai,
        pyfunc,
        pyspark,
        pytorch,
        rfunc,
        semantic_kernel,
        sentence_transformers,
        shap,
        sklearn,
        smolagents,
        spacy,
        spark,
        statsmodels,
        strands,
        tensorflow,
        transformers,
        xgboost,
    )

_install_sensitive_query_param_filter()

if MLFLOW_CONFIGURE_LOGGING.get() is True:
    _configure_mlflow_loggers(root_module_name=__name__)

# Core modules required for mlflow-tracing
from mlflow.tracing.assessment import (
    delete_assessment,
    get_assessment,
    log_assessment,
    log_expectation,
    log_feedback,
    log_issue,
    override_feedback,
    update_assessment,
)
from mlflow.tracing.context import context
from mlflow.tracing.fluent import (
    add_trace,
    delete_trace_tag,
    get_active_trace_id,
    get_current_active_span,
    get_last_active_trace_id,
    get_trace,
    log_trace,
    search_sessions,
    search_traces,
    set_trace_tag,
    start_span,
    start_span_no_context,
    trace,
    update_current_trace,
)
from mlflow.tracking import (
    get_tracking_uri,
    is_tracking_uri_set,
    set_tracking_uri,
)
from mlflow.tracking.fluent import active_run, flush_trace_async_logging, set_experiment

# These are minimal set of APIs to be exposed via `mlflow-tracing` package.
# APIs listed here must not depend on dependencies that are not part of `mlflow-tracing` package.
__all__ = [
    "MlflowException",
    # Minimal tracking APIs required for tracing core functionality
    "set_experiment",
    "set_tracking_uri",
    "get_tracking_uri",
    "is_tracking_uri_set",
    # NB: Tracing SDK doesn't support using Runs, however, active_run is used heavily within
    # the autologging code base.
    "active_run",
    # Tracing APIs
    "add_trace",
    "context",
    "delete_trace_tag",
    "flush_trace_async_logging",
    "get_active_trace_id",
    "get_current_active_span",
    "get_last_active_trace_id",
    "get_trace",
    "log_trace",
    "search_sessions",
    "search_traces",
    "set_trace_tag",
    "start_span",
    "start_span_no_context",
    "trace",
    "update_current_trace",
    # Assessment APIs
    "get_assessment",
    "delete_assessment",
    "log_assessment",
    "update_assessment",
    "log_expectation",
    "log_feedback",
    "log_issue",
    "override_feedback",
]

# Only import these modules when mlflow or mlflow-skinny is installed i.e. not importing them
# when only mlflow-tracing is installed.
if not IS_TRACING_SDK_ONLY:
    from mlflow.client import MlflowClient

    # For backward compatibility, we expose the following functions and classes at the top level in
    # addition to `mlflow.config`.
    from mlflow.config import (
        disable_system_metrics_logging,
        enable_system_metrics_logging,
        get_registry_uri,
        set_registry_uri,
        set_system_metrics_node_id,
        set_system_metrics_samples_before_logging,
        set_system_metrics_sampling_interval,
    )
    from mlflow.models.evaluation.deprecated import evaluate
    from mlflow.models.evaluation.validation import validate_evaluation_results
    from mlflow.projects import run
    from mlflow.pytest import test
    from mlflow.tracking._model_registry.fluent import (
        # TODO: Prompt Registry APIs are moved to the `mlflow.genai` namespace and direct
        # imports from mlflow will be deprecated in the future.
        delete_prompt_alias,
        load_prompt,
        register_model,
        register_prompt,
        search_model_versions,
        search_prompts,
        search_registered_models,
        set_model_version_tag,
        set_prompt_alias,
    )
    from mlflow.tracking._workspace.fluent import (
        create_workspace,
        delete_workspace,
        get_workspace,
        list_workspaces,
        set_workspace,
        update_workspace,
    )
    from mlflow.tracking.fluent import (
        ActiveModel,
        ActiveRun,
        autolog,
        clear_active_model,
        create_experiment,
        create_external_model,
        delete_experiment,
        delete_experiment_tag,
        delete_logged_model_tag,
        delete_run,
        delete_tag,
        end_run,
        finalize_logged_model,
        flush_artifact_async_logging,
        flush_async_logging,
        get_active_model_id,
        get_artifact_uri,
        get_experiment,
        get_experiment_by_name,
        get_logged_model,
        get_parent_run,
        get_run,
        import_checkpoints,
        initialize_logged_model,
        last_active_run,
        last_logged_model,
        load_table,
        log_artifact,
        log_artifacts,
        log_dict,
        log_figure,
        log_image,
        log_input,
        log_inputs,
        log_metric,
        log_metrics,
        log_model_params,
        log_outputs,
        log_param,
        log_params,
        log_stream,
        log_table,
        log_text,
        search_experiments,
        search_logged_models,
        search_runs,
        set_active_model,
        set_experiment_tag,
        set_experiment_tags,
        set_logged_model_tags,
        set_tag,
        set_tags,
        start_run,
    )
    from mlflow.tracking.multimedia import Image
    from mlflow.utils.async_logging.run_operations import RunOperations  # noqa: F401
    from mlflow.utils.credentials import login
    from mlflow.utils.doctor import doctor

    __all__ += [
        "ActiveRun",
        "ActiveModel",
        "MlflowClient",
        "MlflowException",
        "autolog",
        "clear_active_model",
        "create_experiment",
        "create_external_model",
        "create_workspace",
        "delete_experiment",
        "delete_workspace",
        "delete_run",
        "delete_tag",
        "disable_system_metrics_logging",
        "doctor",
        "enable_system_metrics_logging",
        "end_run",
        "evaluate",
        "finalize_logged_model",
        "flush_async_logging",
        "flush_artifact_async_logging",
        "get_active_model_id",
        "get_artifact_uri",
        "get_experiment",
        "get_experiment_by_name",
        "import_checkpoints",
        "get_logged_model",
        "get_workspace",
        "get_parent_run",
        "get_registry_uri",
        "get_run",
        "initialize_logged_model",
        "last_active_run",
        "last_logged_model",
        "load_table",
        "log_artifact",
        "log_artifacts",
        "log_dict",
        "log_figure",
        "log_image",
        "log_input",
        "log_inputs",
        "log_model_params",
        "log_outputs",
        "log_metric",
        "log_metrics",
        "log_param",
        "log_params",
        "log_stream",
        "log_table",
        "log_text",
        "login",
        "pyfunc",
        "register_model",
        "run",
        "search_experiments",
        "search_logged_models",
        "search_model_versions",
        "search_registered_models",
        "list_workspaces",
        "search_runs",
        "search_prompts",
        "set_active_model",
        "set_experiment_tag",
        "set_experiment_tags",
        "delete_experiment_tag",
        "set_model_version_tag",
        "set_registry_uri",
        "set_system_metrics_node_id",
        "set_system_metrics_samples_before_logging",
        "set_system_metrics_sampling_interval",
        "set_tag",
        "set_tags",
        "set_workspace",
        "start_run",
        "test",
        "validate_evaluation_results",
        "Image",
        # Prompt Registry APIs
        # TODO: Prompt Registry APIs are moved to the `mlflow.genai` namespace and direct
        # imports from mlflow will be deprecated in the future.
        "load_prompt",
        "register_prompt",
        "set_prompt_alias",
        "delete_prompt_alias",
        "set_logged_model_tags",
        "delete_logged_model_tag",
        "update_workspace",
    ]


# `mlflow.gateway` depends on optional dependencies such as pydantic, psutil, and has version
# restrictions for dependencies. Importing this module fails if they are not installed or
# if invalid versions of these required packages are installed.
with contextlib.suppress(Exception):
    from mlflow import gateway  # noqa: F401

    __all__.append("gateway")

from mlflow.telemetry import set_telemetry_client

set_telemetry_client()