import inspect
import json
import logging
import os

import cloudpickle

from mlflow.dspy.save import (
    _DSPY_SETTINGS_FILE_NAME,
    _MODEL_CONFIG_FILE_NAME,
    _MODEL_DATA_PATH,
)
from mlflow.dspy.wrapper import DspyChatModelWrapper, DspyModelWrapper
from mlflow.environment_variables import MLFLOW_ALLOW_PICKLE_DESERIALIZATION
from mlflow.exceptions import MlflowException
from mlflow.models import Model
from mlflow.models.dependencies_schemas import _get_dependencies_schema_from_model
from mlflow.models.model import _update_active_model_id_based_on_mlflow_model
from mlflow.tracing.provider import trace_disabled
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.utils.databricks_utils import (
    is_in_databricks_model_serving_environment,
    is_in_databricks_runtime,
)
from mlflow.utils.model_utils import (
    _add_code_from_conf_to_system_path,
    _get_flavor_configuration,
)

_DEFAULT_MODEL_PATH = "data/model.pkl"
_logger = logging.getLogger(__name__)


def _set_dependency_schema_to_tracer(model_path, callbacks):
    """
    Set dependency schemas from the saved model metadata to the tracer
    to propagate it to inference traces.
    """
    from mlflow.dspy.callback import MlflowCallback

    tracer = next((cb for cb in callbacks if isinstance(cb, MlflowCallback)), None)
    if tracer is None:
        return

    model = Model.load(model_path)
    tracer.set_dependencies_schema(_get_dependencies_schema_from_model(model))


def _load_model(model_uri, dst_path=None):
    import dspy

    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
    mlflow_model = Model.load(local_model_path)
    flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name="dspy")

    model_path = flavor_conf.get("model_path", _DEFAULT_MODEL_PATH)
    task = flavor_conf.get("inference_task")

    allow_pickle = (
        MLFLOW_ALLOW_PICKLE_DESERIALIZATION.get()
        or is_in_databricks_runtime()
        or is_in_databricks_model_serving_environment()
    )

    # Raise BEFORE mutating sys.path so a denied load has no global side effects.
    if model_path.endswith(".pkl") and not allow_pickle:
        raise MlflowException(
            "Deserializing model using pickle is disallowed, but this model is saved "
            "in pickle format. To address this issue, you need to set environment variable "
            "'MLFLOW_ALLOW_PICKLE_DESERIALIZATION' to 'true', or save the model with "
            "'use_dspy_model_save=True' like "
            "`mlflow.dspy.save_model(model, path, use_dspy_model_save=True)`."
        )

    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)

    if model_path.endswith(".pkl"):
        with open(os.path.join(local_model_path, model_path), "rb") as f:
            loaded_wrapper = cloudpickle.load(f)
    else:
        try:
            model = dspy.load(os.path.join(local_model_path, model_path), allow_pickle=allow_pickle)
        except Exception as e:
            if not allow_pickle:
                raise MlflowException(
                    f"Failed to load DSPy model: {e}. Note: the environment variable "
                    "'MLFLOW_ALLOW_PICKLE_DESERIALIZATION' is currently set to 'false', "
                    "which disables pickle-based deserialization. If the failure above "
                    "is due to disabled pickle deserialization, set "
                    "'MLFLOW_ALLOW_PICKLE_DESERIALIZATION' to 'true' to allow loading "
                    "pickle-based models."
                ) from e
            raise

        settings_path = os.path.join(local_model_path, _MODEL_DATA_PATH, _DSPY_SETTINGS_FILE_NAME)
        if "allow_pickle" in inspect.signature(dspy.load_settings).parameters:
            dspy_settings = dspy.load_settings(settings_path, allow_pickle=allow_pickle)
        else:
            dspy_settings = dspy.load_settings(settings_path)

        model_config_file = os.path.join(
            local_model_path, _MODEL_DATA_PATH, _MODEL_CONFIG_FILE_NAME
        )
        if os.path.exists(model_config_file):
            with open(model_config_file) as f:
                model_config = json.load(f)
        else:
            model_config = None

        if task == "llm/v1/chat":
            loaded_wrapper = DspyChatModelWrapper(model, dspy_settings, model_config)
        else:
            loaded_wrapper = DspyModelWrapper(model, dspy_settings, model_config)

    _set_dependency_schema_to_tracer(local_model_path, loaded_wrapper.dspy_settings["callbacks"])
    _update_active_model_id_based_on_mlflow_model(mlflow_model)
    return loaded_wrapper


@trace_disabled  # Suppress traces for internal calls while loading model
def load_model(model_uri, dst_path=None):
    """
    Load a Dspy model from a run.

    This function will also set the global dspy settings `dspy.settings` by the saved settings.

    Args:
        model_uri: The location, in URI format, of the MLflow model. For example:

            - ``/Users/me/path/to/local/model``
            - ``relative/path/to/local/model``
            - ``s3://my_bucket/path/to/model``
            - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
            - ``mlflow-artifacts:/path/to/model``

            For more information about supported URI schemes, see
            `Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
            artifact-locations>`_.
        dst_path: The local filesystem path to utilize for downloading the model artifact.
            This directory must already exist if provided. If unspecified, a local output
            path will be created.

    Returns:
        An `dspy.module` instance, representing the dspy model.
    """
    import dspy

    wrapper = _load_model(model_uri, dst_path)

    # Set the global dspy settings for reproducing the model's behavior when the model is
    # loaded via `mlflow.dspy.load_model`. Note that for the model to be loaded as pyfunc,
    # settings will be set in the wrapper's `predict` method via local context to avoid the
    # "dspy.settings can only be changed by the thread that initially configured it" error
    # in Databricks model serving.
    dspy.settings.configure(**wrapper.dspy_settings)

    return wrapper.model


def _load_pyfunc(path):
    return _load_model(path)