"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 5th Dec., 2024.

    OBJECTIVE:

        To provide a structure to normalize input to and output from a standardized LLM wrapper.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


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


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

# For making data behaviour_models:
from pydantic import BaseModel, Field, field_validator, PastDatetime, AwareDatetime
from typing import Optional, Literal, Union, List, Any

# My utils:
from utils_v2.string import regex
from utils_v2.date_time import date_time

# To work with date and time:
import datetime


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


# RegEx Patterns:
REGEX_SESSION_TOKEN = r"^[a-f0-9]{8}-[a-f0-9]{4}-[1-5][a-f0-9]{3}-[89ab][a-f0-9]{3}-[a-f0-9]{12}$"


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


# --- Nothing Yet


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


class LLMInputMessage(BaseModel):

    role: Literal["system", "ai", "human"] = Field(
        description = "the role of this message",
        frozen = True
    )

    content: str = Field(
        description = "the message sent by the 'role'",
        frozen = True
    )

    # ┏┓    ┏•
    # ┃ ┏┓┏┓╋┓┏┓
    # ┗┛┗┛┛┗┛┗┗┫
    #          ┛

    class Config:
        extra = "forbid"


# ---------------------------------------------------------------------------------------------------------------------


class LLMInput(BaseModel):

    messages: List[LLMInputMessage]

    # ┏┓    ┏•
    # ┃ ┏┓┏┓╋┓┏┓
    # ┗┛┗┛┛┗┛┗┗┫
    #          ┛

    class Config:
        extra = "forbid"

    # ┓┏  ┓• ┓   •
    # ┃┃┏┓┃┓┏┫┏┓╋┓┏┓┏┓
    # ┗┛┗┻┗┗┗┻┗┻┗┗┗┛┛┗

    @field_validator("messages")
    def validate_messages(cls, value):

        # Maintain counter(s):
        system_message_index = -1
        system_message_count = 0

        # Loop through the messages and check them:
        for index, message in enumerate(value):

            # For 'system' messages:
            if message.role == "system":
                system_message_index = index
                system_message_count += 1

        # Verify that there is AT MOST ONE 'system' message,
        # and verify that the 'system' message is the first message:
        if system_message_count > 1: raise ValueError(f"there can be at most 1 'system' message, found {system_message_count}")
        if system_message_index > 0: raise ValueError(f"'system' message must always be at index 0, found it at index {system_message_index}")

        # Done here:
        return value


# ---------------------------------------------------------------------------------------------------------------------


class LLMUsageTokens(BaseModel):

    input: int = Field(
        description = "how many tokens were given in the input",
        frozen = True
    )

    output: int = Field(
        description = "how many tokens were generated as the output",
        frozen = True
    )

    total: int = Field(
        description = "the sum of the input and output tokens",
        frozen = True
    )

    # ┏┓    ┏•
    # ┃ ┏┓┏┓╋┓┏┓
    # ┗┛┗┛┛┗┛┗┗┫
    #          ┛

    class Config:
        extra = "forbid"


# ---------------------------------------------------------------------------------------------------------------------


class LLMOutput(BaseModel):

    ts: AwareDatetime = Field(
        description = "the time at which the llm was invoked",
        default_factory = date_time.get_current_utc_date_time,
        frozen = True
    )

    messages: List[LLMInputMessage] = Field(
        description = "the messages that came in that invoked the llm",
        frozen = True
    )

    output: str | None = Field(
        description = "what the llm generated",
        default = None,
        frozen = True
    )

    client: Literal["openai"] = Field(
        description = "the co./brand that was used to use an llm",
        frozen = True
    )

    model: str = Field(
        description = "to know which model used in the process",
        frozen = True
    )

    tokens: LLMUsageTokens = Field(
        description = "to know how many tokens were used in the process",
        default = LLMUsageTokens(input = 0, output = 0, total = 0),
        frozen = True
    )

    invocationId: Any | None = Field(
        description = "the id of the document that notes this invocation; useful for reconciliation",
        frozen = False,
        default = None
    )

    # ┏┓    ┏•
    # ┃ ┏┓┏┓╋┓┏┓
    # ┗┛┗┛┛┗┛┗┗┫
    #          ┛

    class Config:
        extra = "forbid"


# ---------------------------------------------------------------------------------------------------------------------


class LLMRequestHeaders(BaseModel):

    sessionToken: str = Field(
        description = "the session token of the user who is requesting the service",
        pattern = REGEX_SESSION_TOKEN,
        frozen = True,
        alias = "X-Session-Token"
    )

    # ┏┓    ┏•
    # ┃ ┏┓┏┓╋┓┏┓
    # ┗┛┗┛┛┗┛┗┗┫
    #          ┛

    class Config:
        extra = "allow"

    def model_dump(self, *args, **kwargs):
        return super().model_dump(*args, by_alias = True, **kwargs)


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


if __name__ == "__main__":

    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..."
        }
    ]

    llm_input = LLMInput(messages = llm_messages)
    print(llm_input)