"""

    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.string import json
from utils_v2.date_time import date_time
from utils_v2.database.async_mysql_v2 import AsyncMySQL
from utils_v2.database.async_mongo_v2 import AsyncMongo

# Base model:
from models.behaviour.base import BaseModel

# Data Models:
from models.data.ai.llm import LLMInput, LLMOutput, LLMUsageTokens

# To work with LLMs:
from langchain_openai import ChatOpenAI

# To work with MongoDB:
from bson import ObjectId

# To work with datatypes:
from typing import Literal

# To make deep-copies:
import copy


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


# --- Nothing Yet


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


# --- Nothing Yet


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


# --- Nothing Yet


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


class LLMOpenAI(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: dict,
            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}
        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
        )

        # Done here:
        return llm_response


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


if __name__ == "__main__":

    pass

    # import asyncio
    #
    # llm_messages = [
    #     {
    #         "role": "system",
    #         "content": "You are an office assistant."
    #     },
    #     {
    #         "role": "ai",
    #         "content": "Hello, sir. How may I help you today?"
    #     },
    #     {
    #         "role": "human",
    #         "content": "Please summarize this mail for me..."
    #     }
    # ]
    #
    # my_llm = LLMOpenAI(
    #     llm_creds = {
    #         "model": "gpt-4o-mini",
    #         "openai_api_key": "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"
    #     }
    # )
    #
    # async def main():
    #
    #     llm_response = await my_llm.invoke(llm_input = LLMInput(messages = llm_messages))
    #     print("LLM RESPONSE:", llm_response.model_dump_json(indent = 4))
    #
    # asyncio.run(main())