"""Wrapper around IBM GENAI APIs for use in langchain"""
import logging
from typing import Any, List, Mapping, Optional

from pydantic import BaseModel, Extra

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 ibm-generative-ai[langchain] extension.")

from genai import Credentials, Model
from genai.schemas import GenerateParams

logger = logging.getLogger(__name__)

__all__ = ["LangChainInterface"]


class LangChainInterface(LLM, BaseModel):
    """
    Wrapper around IBM GENAI models.
    To use, you should have the ``genai`` python package installed
    and initialize the credentials attribute of this class with
    an instance of ``genai.Credentials``. Model specific parameters
    can be passed through to the constructor using the ``params``
    parameter, which is an instance of GenerateParams.
    Example:
        .. code-block:: python
            llm = LangChainInterface(model="google/flan-ul2", credentials=creds)
    """

    credentials: Credentials = None
    model: Optional[str] = None
    params: Optional[GenerateParams] = None

    class Config:
        """Configuration for this pydantic object."""

        extra = Extra.forbid

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        _params = self.params or GenerateParams()
        return {
            **{"model": self.model},
            **{"params": _params},
        }

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

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        """Call the IBM GENAI's 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 = LangChainInterface(model_id="google/flan-ul2", credentials=creds)
                response = llm("What is a molecule")
        """
        params = self.params or GenerateParams()
        model = Model(model=self.model, params=params, credentials=self.credentials)
        text = model.generate(prompts=[prompt])[0].generated_text
        logger.info("Output of GENAI call: {}".format(text))
        if stop is not None:
            text = enforce_stop_tokens(text, stop)
        return text
