"""Utility functions for mlflow.langchain."""

import functools
import importlib
import json
import logging
import os
import shutil
import types
from functools import lru_cache
from importlib.util import find_spec
from typing import Any, Callable, NamedTuple

import cloudpickle
import yaml
from packaging.version import Version

import mlflow
from mlflow.models.utils import _validate_and_get_model_code_path
from mlflow.utils.class_utils import _get_class_from_string

_AGENT_PRIMITIVES_FILE_NAME = "agent_primitive_args.json"
_AGENT_PRIMITIVES_DATA_KEY = "agent_primitive_data"
_AGENT_DATA_FILE_NAME = "agent.yaml"
_AGENT_DATA_KEY = "agent_data"
_TOOLS_DATA_FILE_NAME = "tools.pkl"
_TOOLS_DATA_KEY = "tools_data"
_LOADER_FN_FILE_NAME = "loader_fn.pkl"
_LOADER_FN_KEY = "loader_fn"
_LOADER_ARG_KEY = "loader_arg"
_PERSIST_DIR_NAME = "persist_dir_data"
_PERSIST_DIR_KEY = "persist_dir"
_MODEL_DATA_YAML_FILE_NAME = "model.yaml"
_MODEL_DATA_PKL_FILE_NAME = "model.pkl"
_MODEL_DATA_FOLDER_NAME = "model"
_MODEL_DATA_KEY = "model_data"
_MODEL_TYPE_KEY = "model_type"
_RUNNABLE_LOAD_KEY = "runnable_load"
_BASE_LOAD_KEY = "base_load"
_CONFIG_LOAD_KEY = "config_load"
_PICKLE_LOAD_KEY = "pickle_load"
_MODEL_LOAD_KEY = "model_load"
_UNSUPPORTED_MODEL_WARNING_MESSAGE = (
    "MLflow does not guarantee support for Chains outside of the subclasses of LLMChain, found %s"
)
_UNSUPPORTED_LLM_WARNING_MESSAGE = (
    "MLflow does not guarantee support for LLMs outside of HuggingFacePipeline and OpenAI, found %s"
)


try:
    import langchain_community

    # Since langchain-community 0.0.27, saving or loading a module that relies on the pickle
    # deserialization requires passing `allow_dangerous_deserialization=True`.
    IS_PICKLE_SERIALIZATION_RESTRICTED = Version(langchain_community.__version__) >= Version(
        "0.0.27"
    )
except ImportError:
    IS_PICKLE_SERIALIZATION_RESTRICTED = False

logger = logging.getLogger(__name__)


@lru_cache
def base_lc_types():
    """
    Get base LangChain types (Chain, AgentExecutor, BaseRetriever).

    Note: AgentExecutor was removed in langchain 1.0.0. Use LangGraph instead.
    """
    from mlflow.langchain._compat import (
        import_base_retriever,
        try_import_agent_executor,
        try_import_chain,
    )

    types = []

    if chain_cls := try_import_chain():
        types.append(chain_cls)

    if agent_executor_cls := try_import_agent_executor():
        types.append(agent_executor_cls)

    types.append(import_base_retriever())

    return tuple(types)


@lru_cache
def picklable_runnable_types():
    """
    Runnable types that can be pickled and unpickled by cloudpickle.
    """
    from mlflow.langchain._compat import (
        import_chat_prompt_template,
        import_runnable_lambda,
        import_runnable_passthrough,
        try_import_simple_chat_model,
    )

    types = [
        import_chat_prompt_template(),
        import_runnable_passthrough(),
        import_runnable_lambda(),
    ]

    if simple_chat_model := try_import_simple_chat_model():
        types.insert(0, simple_chat_model)

    return tuple(types)


@lru_cache
def lc_runnable_with_steps_types():
    from mlflow.langchain._compat import import_runnable_parallel, import_runnable_sequence

    return (import_runnable_parallel(), import_runnable_sequence())


def lc_runnable_assign_types():
    from mlflow.langchain._compat import import_runnable_assign

    return (import_runnable_assign(),)


def lc_runnable_branch_types():
    from mlflow.langchain._compat import import_runnable_branch

    return (import_runnable_branch(),)


def lc_runnable_binding_types():
    from mlflow.langchain._compat import import_runnable_binding

    return (import_runnable_binding(),)


def lc_runnables_types():
    return (
        picklable_runnable_types()
        + lc_runnable_with_steps_types()
        + lc_runnable_branch_types()
        + lc_runnable_assign_types()
        + lc_runnable_binding_types()
    )


def langgraph_types():
    try:
        from langgraph.graph.state import CompiledStateGraph

        return (CompiledStateGraph,)
    except ImportError:
        return ()


def supported_lc_types():
    return base_lc_types() + lc_runnables_types() + langgraph_types()


# Wrapping as a function to avoid calling supported_lc_types() at import time
def get_unsupported_model_message(model_type):
    return (
        "MLflow langchain flavor only supports subclasses of "
        f"{supported_lc_types()}, found {model_type}."
    )


@lru_cache
def custom_type_to_loader_dict():
    # helper function to load output_parsers from config
    def _load_output_parser(config: dict[str, Any]) -> Any:
        """Load output parser."""
        from mlflow.langchain._compat import import_str_output_parser

        output_parser_type = config.pop("_type", None)
        if output_parser_type == "default":
            return import_str_output_parser()(**config)
        else:
            raise ValueError(f"Unsupported output parser {output_parser_type}")

    return {"default": _load_output_parser}


class _SpecialChainInfo(NamedTuple):
    loader_arg: str


def _get_special_chain_info_or_none(chain):
    for (
        special_chain_class,
        loader_arg,
    ) in _get_map_of_special_chain_class_to_loader_arg().items():
        if isinstance(chain, special_chain_class):
            return _SpecialChainInfo(loader_arg=loader_arg)


@lru_cache
def _get_map_of_special_chain_class_to_loader_arg():
    class_name_to_loader_arg = {
        "langchain.chains.RetrievalQA": "retriever",
        "langchain.chains.APIChain": "requests_wrapper",
        "langchain.chains.HypotheticalDocumentEmbedder": "embeddings",
    }
    # SQLDatabaseChain is in langchain_experimental (since version 0.0.247+)
    if find_spec("langchain_experimental"):
        # Add this entry only if langchain_experimental is installed
        class_name_to_loader_arg["langchain_experimental.sql.SQLDatabaseChain"] = "database"

    class_to_loader_arg = {}
    try:
        from mlflow.langchain.retriever_chain import _RetrieverChain

        class_to_loader_arg[_RetrieverChain] = "retriever"
    except ImportError:
        pass

    for class_name, loader_arg in class_name_to_loader_arg.items():
        try:
            cls = _get_class_from_string(class_name)
            class_to_loader_arg[cls] = loader_arg
        except Exception:
            logger.warning(
                "Unexpected import failure for class '%s'. Please file an issue at"
                " https://github.com/mlflow/mlflow/issues/.",
                class_name,
                exc_info=True,
            )

    return class_to_loader_arg


@lru_cache
def _get_supported_llms():
    supported_llms = set()

    def try_adding_llm(module, class_name):
        if cls := getattr(module, class_name, None):
            supported_llms.add(cls)

    def safe_import_and_add(module_name, class_name):
        """Add conditional support for `partner` and `community` APIs in langchain"""
        try:
            module = importlib.import_module(module_name)
            try_adding_llm(module, class_name)
        except ImportError:
            pass

    safe_import_and_add("langchain.llms.openai", "OpenAI")
    # HuggingFacePipeline is moved to langchain_huggingface since langchain 0.2.0
    safe_import_and_add("langchain.llms", "HuggingFacePipeline")
    safe_import_and_add("langchain.langchain_huggingface", "HuggingFacePipeline")
    safe_import_and_add("langchain_openai", "OpenAI")
    safe_import_and_add("langchain_databricks", "ChatDatabricks")
    safe_import_and_add("databricks_langchain", "ChatDatabricks")

    for llm_name in ["Databricks", "Mlflow"]:
        safe_import_and_add("langchain.llms", llm_name)

    for chat_model_name in [
        "ChatDatabricks",
        "ChatMlflow",
        "ChatOpenAI",
        "AzureChatOpenAI",
    ]:
        safe_import_and_add("langchain.chat_models", chat_model_name)

    return supported_llms


def _agent_executor_contains_unsupported_llm(lc_model, _SUPPORTED_LLMS):
    from mlflow.langchain._compat import try_import_agent_executor

    agent_executor_cls = try_import_agent_executor()
    if agent_executor_cls is None:
        return False

    return (
        isinstance(lc_model, agent_executor_cls)
        # 'RunnableMultiActionAgent' object has no attribute 'llm_chain'
        and hasattr(lc_model.agent, "llm_chain")
        and not any(
            isinstance(lc_model.agent.llm_chain.llm, supported_llm)
            for supported_llm in _SUPPORTED_LLMS
        )
    )


# temp_dir is only required when lc_model could be a file path
def _validate_and_prepare_lc_model_or_path(lc_model, loader_fn, temp_dir=None):
    if isinstance(lc_model, str):
        return _validate_and_get_model_code_path(lc_model, temp_dir)

    if not isinstance(lc_model, supported_lc_types()):
        raise mlflow.MlflowException.invalid_parameter_value(
            get_unsupported_model_message(type(lc_model).__name__)
        )

    _SUPPORTED_LLMS = _get_supported_llms()

    from mlflow.langchain._compat import try_import_llm_chain

    llm_chain_cls = try_import_llm_chain()
    if (
        llm_chain_cls
        and isinstance(lc_model, llm_chain_cls)
        and not any(isinstance(lc_model.llm, supported_llm) for supported_llm in _SUPPORTED_LLMS)
    ):
        logger.warning(
            _UNSUPPORTED_LLM_WARNING_MESSAGE,
            type(lc_model.llm).__name__,
        )

    if _agent_executor_contains_unsupported_llm(lc_model, _SUPPORTED_LLMS):
        logger.warning(
            _UNSUPPORTED_LLM_WARNING_MESSAGE,
            type(lc_model.agent.llm_chain.llm).__name__,
        )

    if special_chain_info := _get_special_chain_info_or_none(lc_model):
        if loader_fn is None:
            raise mlflow.MlflowException.invalid_parameter_value(
                f"For {type(lc_model).__name__} models, a `loader_fn` must be provided."
            )
        if not isinstance(loader_fn, types.FunctionType):
            raise mlflow.MlflowException.invalid_parameter_value(
                "The `loader_fn` must be a function that returns a {loader_arg}.".format(
                    loader_arg=special_chain_info.loader_arg
                )
            )

    # If lc_model is a retriever, wrap it in a _RetrieverChain
    from mlflow.langchain._compat import import_base_retriever

    BaseRetriever = import_base_retriever()
    if isinstance(lc_model, BaseRetriever):
        try:
            from mlflow.langchain.retriever_chain import _RetrieverChain
        except ImportError:
            raise mlflow.MlflowException.invalid_parameter_value(
                "_RetrieverChain is not available. It requires langchain<1.0.0. "
                "For langchain>=1.0.0, please use LangGraph instead."
            )

        if loader_fn is None:
            raise mlflow.MlflowException.invalid_parameter_value(
                f"For {type(lc_model).__name__} models, a `loader_fn` must be provided."
            )
        if not isinstance(loader_fn, types.FunctionType):
            raise mlflow.MlflowException.invalid_parameter_value(
                "The `loader_fn` must be a function that returns a retriever."
            )
        lc_model = _RetrieverChain(retriever=lc_model)

    return lc_model


def _save_base_lcs(model, path, loader_fn=None, persist_dir=None):
    from mlflow.langchain._compat import (
        try_import_agent_executor,
        try_import_base_chat_model,
        try_import_chain,
        try_import_llm_chain,
    )

    AgentExecutor = try_import_agent_executor()
    Chain = try_import_chain()
    LLMChain = try_import_llm_chain()
    BaseChatModel = try_import_base_chat_model()

    model_data_path = os.path.join(path, _MODEL_DATA_YAML_FILE_NAME)
    model_data_kwargs = {
        _MODEL_DATA_KEY: _MODEL_DATA_YAML_FILE_NAME,
        _MODEL_LOAD_KEY: _BASE_LOAD_KEY,
    }

    is_llm_chain = LLMChain and isinstance(model, LLMChain)
    is_base_chat_model = BaseChatModel and isinstance(model, BaseChatModel)

    if is_llm_chain or is_base_chat_model:
        model.save(model_data_path)
    elif AgentExecutor and isinstance(model, AgentExecutor):
        if model.agent and getattr(model.agent, "llm_chain", None):
            model.agent.llm_chain.save(model_data_path)

        if model.agent:
            agent_data_path = os.path.join(path, _AGENT_DATA_FILE_NAME)
            model.save_agent(agent_data_path)
            model_data_kwargs[_AGENT_DATA_KEY] = _AGENT_DATA_FILE_NAME

        if model.tools:
            tools_data_path = os.path.join(path, _TOOLS_DATA_FILE_NAME)
            try:
                with open(tools_data_path, "wb") as f:
                    cloudpickle.dump(model.tools, f)
            except Exception as e:
                raise mlflow.MlflowException(
                    "Error when attempting to pickle the AgentExecutor tools. "
                    "This model likely does not support serialization."
                ) from e
            model_data_kwargs[_TOOLS_DATA_KEY] = _TOOLS_DATA_FILE_NAME
        else:
            raise mlflow.MlflowException.invalid_parameter_value(
                "For initializing the AgentExecutor, tools must be provided."
            )

        key_to_ignore = ["llm_chain", "agent", "tools", "callback_manager"]
        temp_dict = {k: v for k, v in model.__dict__.items() if k not in key_to_ignore}

        agent_primitive_path = os.path.join(path, _AGENT_PRIMITIVES_FILE_NAME)
        with open(agent_primitive_path, "w") as config_file:
            json.dump(temp_dict, config_file, indent=4)

        model_data_kwargs[_AGENT_PRIMITIVES_DATA_KEY] = _AGENT_PRIMITIVES_FILE_NAME

    elif special_chain_info := _get_special_chain_info_or_none(model):
        # Save loader_fn by pickling
        loader_fn_path = os.path.join(path, _LOADER_FN_FILE_NAME)
        with open(loader_fn_path, "wb") as f:
            cloudpickle.dump(loader_fn, f)
        model_data_kwargs[_LOADER_FN_KEY] = _LOADER_FN_FILE_NAME
        model_data_kwargs[_LOADER_ARG_KEY] = special_chain_info.loader_arg

        if persist_dir is not None:
            if os.path.exists(persist_dir):
                # Save persist_dir by copying into subdir _PERSIST_DIR_NAME
                persist_dir_data_path = os.path.join(path, _PERSIST_DIR_NAME)
                shutil.copytree(persist_dir, persist_dir_data_path)
                model_data_kwargs[_PERSIST_DIR_KEY] = _PERSIST_DIR_NAME
            else:
                raise mlflow.MlflowException.invalid_parameter_value(
                    "The directory provided for persist_dir does not exist."
                )

        # Save model
        model.save(model_data_path)
    elif Chain and isinstance(model, Chain):
        logger.warning(get_unsupported_model_message(type(model).__name__))
        model.save(model_data_path)
    else:
        raise mlflow.MlflowException.invalid_parameter_value(
            get_unsupported_model_message(type(model).__name__)
        )

    return model_data_kwargs


def _load_from_pickle(path):
    with open(path, "rb") as f:
        return cloudpickle.load(f)


def _load_from_json(path):
    with open(path) as f:
        return json.load(f)


def _load_from_yaml(path):
    with open(path) as f:
        return yaml.safe_load(f)


def _get_path_by_key(root_path, key, conf):
    key_path = conf.get(key)
    return os.path.join(root_path, key_path) if key_path else None


def _patch_loader(loader_func: Callable[..., Any]) -> Callable[..., Any]:
    """
    Patch LangChain loader function like load_chain() to handle pickle deserialization.

    Since langchain-community 0.0.27, loading a module that relies on the pickle deserialization
    requires the `allow_dangerous_deserialization` flag to be set to True, for security reasons.

    Args:
        loader_func: The LangChain loader function to be patched e.g. load_chain().

    Returns:
        The patched loader function.
    """
    if not IS_PICKLE_SERIALIZATION_RESTRICTED:
        return loader_func

    # For LangChain >= 0.3.0, we can pass `allow_dangerous_deserialization` flag
    # via the loader APIs. Since the model is serialized by the user (or someone who has
    # access to the tracking server), it is safe to set this flag to True.
    def patched_loader(*args, **kwargs):
        return loader_func(*args, **kwargs, allow_dangerous_deserialization=True)

    return patched_loader


def _load_base_lcs(
    local_model_path,
    conf,
):
    lc_model_path = os.path.join(
        local_model_path, conf.get(_MODEL_DATA_KEY, _MODEL_DATA_YAML_FILE_NAME)
    )

    agent_path = _get_path_by_key(local_model_path, _AGENT_DATA_KEY, conf)
    tools_path = _get_path_by_key(local_model_path, _TOOLS_DATA_KEY, conf)
    agent_primitive_path = _get_path_by_key(local_model_path, _AGENT_PRIMITIVES_DATA_KEY, conf)
    loader_fn_path = _get_path_by_key(local_model_path, _LOADER_FN_KEY, conf)
    persist_dir = _get_path_by_key(local_model_path, _PERSIST_DIR_KEY, conf)

    model_type = conf.get(_MODEL_TYPE_KEY)
    loader_arg = conf.get(_LOADER_ARG_KEY)

    load_chain = None
    try:
        from langchain.chains.loading import load_chain
    except ImportError:
        pass

    _RetrieverChain = None
    try:
        from mlflow.langchain.retriever_chain import _RetrieverChain
    except ImportError:
        pass

    if loader_arg is not None:
        if loader_fn_path is None:
            raise mlflow.MlflowException.invalid_parameter_value(
                "Missing file for loader_fn which is required to build the model."
            )
        loader_fn = _load_from_pickle(loader_fn_path)
        kwargs = {loader_arg: loader_fn(persist_dir)}
        if _RetrieverChain and model_type == _RetrieverChain.__name__:
            model = _RetrieverChain.load(lc_model_path, **kwargs).retriever
        else:
            if load_chain is None:
                raise mlflow.MlflowException(
                    "Cannot load model: langchain.chains.loading.load_chain is not available. "
                    "This may be because you're using langchain>=1.0.0. "
                    "Please use a model saved with langchain>=1.0.0."
                )
            model = _patch_loader(load_chain)(lc_model_path, **kwargs)
    elif agent_path is None and tools_path is None:
        if load_chain is None:
            raise mlflow.MlflowException(
                "Cannot load model: langchain.chains.loading.load_chain is not available. "
                "This may be because you're using langchain>=1.0.0. "
                "Please use a model saved with langchain>=1.0.0."
            )
        model = _patch_loader(load_chain)(lc_model_path)
    else:
        try:
            from langchain.agents import initialize_agent
        except ImportError:
            raise mlflow.MlflowException(
                "Cannot load AgentExecutor: langchain.agents.initialize_agent is not available. "
                "AgentExecutor was removed in langchain 1.0.0. Please use LangGraph instead."
            )

        if load_chain is None:
            raise mlflow.MlflowException(
                "Cannot load model: langchain.chains.loading.load_chain is not available. "
                "This may be because you're using langchain>=1.0.0. "
                "Please use a model saved with langchain>=1.0.0."
            )

        llm = _patch_loader(load_chain)(lc_model_path)
        tools = []
        kwargs = {}

        if os.path.exists(tools_path):
            tools = _load_from_pickle(tools_path)
        else:
            raise mlflow.MlflowException(
                "Missing file for tools which is required to build the AgentExecutor object."
            )

        if os.path.exists(agent_primitive_path):
            kwargs = _load_from_json(agent_primitive_path)

        model = initialize_agent(tools=tools, llm=llm, agent_path=agent_path, **kwargs)
    return model


def patch_langchain_type_to_cls_dict(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        def _load_chat_openai():
            from langchain_community.chat_models import ChatOpenAI

            return ChatOpenAI

        def _load_azure_chat_openai():
            from langchain_community.chat_models import AzureChatOpenAI

            return AzureChatOpenAI

        def _load_chat_databricks():
            from databricks_langchain import ChatDatabricks

            return ChatDatabricks

        def _patched_get_type_to_cls_dict(original):
            def _wrapped():
                return {
                    **original(),
                    "openai-chat": _load_chat_openai,
                    "azure-openai-chat": _load_azure_chat_openai,
                    "chat-databricks": _load_chat_databricks,
                }

            return _wrapped

        modules_to_patch = [
            "langchain_databricks",
            "langchain.llms",
            "langchain_community.llms.loading",
        ]
        originals = {}
        for name in modules_to_patch:
            try:
                module = importlib.import_module(name)
                originals[name] = module.get_type_to_cls_dict  # Record original impl for cleanup
            except (ImportError, AttributeError):
                continue
            module.get_type_to_cls_dict = _patched_get_type_to_cls_dict(originals[name])

        try:
            return func(*args, **kwargs)
        finally:
            # Clean up the patch
            for module_name, original_impl in originals.items():
                module = importlib.import_module(module_name)
                module.get_type_to_cls_dict = original_impl

    return wrapper