#  -----------------------------------------------------------------------------------------
#  (C) Copyright IBM Corp. 2023- 2023.
#  https://opensource.org/licenses/BSD-3-Clause
#  -----------------------------------------------------------------------------------------
import logging
from typing import Any, List, Mapping, Optional, Union

try:
    from langchain.llms.base import LLM
    from langchain.llms.utils import enforce_stop_tokens
except ImportError:
    raise ImportError("Could not import langchain: Please install langchain extension.")
from ibm_watson_machine_learning.foundation_models.model import Model

logger = logging.getLogger(__name__)


class WatsonxLLM(LLM):
    """
    Wrapper around IBM watsonx.ai

    Model specific parameters can be passed through to the constructor using the ``params``
    parameter, which is an instance of ``GenTextParamsMetaNames``.

    :param model_id: the type of model to use
    :type model_id: str

    :param credentials: credentials to Watson Machine Learning instance
    :type credentials: dict

    :param params: parameters to use during generate requests
    :type params: dict, optional

    :param project_id: ID of the Watson Studio project
    :type project_id: str, optional

    :param space_id: ID of the Watson Studio space
    :type space_id: str, optional

    :param verify: user can pass as verify one of following:

        - the path to a CA_BUNDLE file
        - the path of directory with certificates of trusted CAs
        - `True` - default path to truststore will be taken
        - `False` - no verification will be made
    :type verify: bool or str, optional

    .. note::
        One of these parameters is required: ['project_id ', 'space_id']

    .. hint::
        You can copy the project_id from Project's Manage tab (Project -> Manage -> General -> Details).

    **Instantiate the WatsonxLLM interface**

    .. code-block:: python

        from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams
        from ibm_watson_machine_learning.foundation_models.extensions.langchain.llm import WatsonxLLM

        generate_params = {
            GenParams.MAX_NEW_TOKENS: 25
        }

        llm = WatsonxLLM(
            model_id = "google/flan-ul2",
            credentials = {
                "apikey": "***",
                "url": "https://us-south.ml.cloud.ibm.com"
            },
            params=generate_params,
            project_id = "*****"
            )

    """

    model_id: str = None
    credentials: dict = None
    params: dict = None
    project_id: str = None
    space_id: str = None
    verify: Union[str, bool] = None
    model: object = None
    llm_type: str = "IBM watsonx.ai"

    def __init__(self,
                 model_id: str,
                 credentials: dict,
                 params: dict = None,
                 project_id: str = None,
                 space_id: str = None,
                 verify: Union[str, bool] = None) -> None:
        super(WatsonxLLM, self).__init__()
        self.model_id = model_id
        self.params = params
        self.project_id = project_id
        self.space_id = space_id
        self.verify = verify
        self.credentials = credentials
        self.model = Model(
            model_id=model_id,
            credentials=credentials,
            params=params,
            project_id=project_id,
            space_id=space_id,
            verify=verify)

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {
            **{"model_id": self.model_id},
            **{"credentials": {
                "url": self.credentials['url'],
                "apikey": "*****"}},
            **{"params": self.params},
            **{"project_id": self.project_id},
            **{"space_id": self.space_id},
            **{"verify": self.verify}
        }

    @property
    def _llm_type(self) -> str:
        """Return type of llm."""
        return self.llm_type

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        """Call the IBM watsonx.ai inference endpoint.
        Args:
            prompt: The prompt to pass into the model.
            stop: Optional list of stop words to use when generating.
        Returns:
            The string generated by the model.
        Example:
            .. code-block:: python
                llm = WatsonxLLM(
                    model_id="google/flan-ul2",
                    credentials={
                        "apikey": "***",
                        "url": "https://us-south.ml.cloud.ibm.com"
                        },
                    project_id="*****"
                    )
                response = llm("What is a molecule")
        """
        text = self.model.generate_text(prompt=prompt)
        logger.info("Output of watsonx.ai call: {}".format(text))
        if stop is not None:
            text = enforce_stop_tokens(text, stop)
        return text
