"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 5th Dec., 2024

    OBJECTIVE:

        To create an interface between OpenAI and our internal system to perform LLM-based activities.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""

#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# My async utils:
from utils_v2.database.async_mongo_v2 import AsyncMongo

# Base model:
from controllers.base import BaseModel

# Data Models:
from models.core.user import CoreUserInfoModel
from models.core.ai.llm import LLMInput, LLMOutput, LLMUsageTokens

# To work with LLMs:
from langchain_openai import ChatOpenAI


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class CoreLLMController(BaseModel):

    AI_USAGE_COLLECTION = "_aiUsage"

    def __init__(
            self,
            llm_creds: dict,
            cache = None,
            alert_url = None,
            http_client = None,
            debug = True,
            debug_prefix = "Model | ",
            debug_only_errors = True
    ):

        """
        This is the model that works with OpenAi's LLM to perform tasks like text completion.
        :param llm_creds: The JSON that holds the credentials to access your OpenAI account. Should have the keys
            'model', and 'openai_api_key'.
        :param cache: The object to use for caching results from database calls.
        :param alert_url: Which URL to call when something goes wrong.
        :param http_client: The instance of an HTTP client to use when trying to send alerts and make other APIs.
        :param debug: Whether, or not, you would like to print debugging messages:
        :param debug_prefix: The prefix to print with the debugging messages.
        :param debug_only_errors: Whether you would like to print only error messages or all messages.
        :return: None.
        """

        # Initialize the parent:
        super().__init__(
            cache = cache,
            alert_url = alert_url,
            http_client = http_client,
            debug = debug,
            debug_prefix = debug_prefix,
            debug_only_errors = debug_only_errors
        )

        # Create the interface to the LLM:
        self.__llm = ChatOpenAI(**llm_creds)

    async def invoke(
            self,
            mongo_conn: AsyncMongo,
            user_info: CoreUserInfoModel,
            llm_input: LLMInput
    ) -> LLMOutput:

        # Format the message as per the format of OpenAI:
        prompt = [
            {
                "role": {"system": "system", "ai": "assistant", "human": "user"}[message.role],
                "content": message.content
            } for message in llm_input.messages
        ]

        # Invoke the AI, and format the response:
        llm_response = await self.__llm.ainvoke(prompt)
        llm_response = LLMOutput(
            messages = llm_input.messages,
            output = llm_response.content,
            client = "openai",
            model = llm_response.response_metadata["model_name"],
            tokens = LLMUsageTokens(
                input = llm_response.usage_metadata["input_tokens"],
                output = llm_response.usage_metadata["output_tokens"],
                total = llm_response.usage_metadata["total_tokens"],
            )
        )

        # Store this into MongoDB:
        mongo_document = {"user": user_info.model_dump()}
        for k, v in llm_response.model_dump().items(): mongo_document[k] = v
        inserted_id = await mongo_conn.insert_one(
            collection = self.AI_USAGE_COLLECTION,
            document = mongo_document
        )
        if inserted_id: llm_response.invocationId = str(inserted_id)

        # Done here:
        return llm_response


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    pass
