from typing import Optional

from openai import OpenAI
from pydantic import BaseModel


class ImageChatHistory:
    def __init__(self):
        self.messages = []

    def clear(self):
        self.messages = []

    def add_system_msg(self, content: str):
        self.messages.append({"role": "system", "content": "<|ENDOFTEXT|>" + content})

    def add_user_msg(self, content: str):
        self.messages.append(
            {"role": "user", "content": [{"type": "text", "text": content}]}
        )

    def add_assistant_msg(self, content: str):
        self.messages.append(
            {"role": "assistant", "content": [{"type": "text", "text": content}]}
        )

    def __repr__(self) -> str:
        strr = "::: Conversation History :::\n"
        for message in self.messages:
            role = message["role"]
            if role == "system":
                strr = strr + role + ":\n" + message["content"] + "\n"
            else:
                content = ""
                for c in message["content"]:
                    if c["type"] == "text":
                        content = content + c["text"] + "\n"
                strr = strr + role + ":\n" + content + "\n"
        return strr

    def __str__(self) -> str:
        return self.__repr__()


def call_model(
    history: ImageChatHistory,
    structured_output: BaseModel | None = None,
    temperature: float = 1.0,  # 1 is openai default, value between 0-2 see also top_p in https://platform.openai.com/docs/api-reference/chat/create
    model: str = "gpt-4o-mini",
    max_response_tokens: int = 4096,  # this is max for gpt-4o as of 'now' (2024-10-06)
    api_key: Optional[
        str
    ] = None,  # if not provided will use env variable OPENAI_API_KEY
):
    client = OpenAI() if api_key is None else OpenAI(api_key=api_key)
    try:
        if structured_output is not None:
            # Use parse for structured output
            res = client.beta.chat.completions.parse(
                model=model,
                messages=history.messages,
                max_tokens=max_response_tokens,
                temperature=temperature,
                response_format=structured_output,
            )
            content = res.choices[0].message.parsed
            return True, content, None
        else:
            # Use create for unstructured output
            res = client.chat.completions.create(
                model=model,
                messages=history.messages,
                max_tokens=max_response_tokens,
                temperature=temperature,
            )
            content = res.choices[0].message.content
            return True, content, None

    except Exception as e:
        # Catch any other exceptions
        print(f"An unexpected error occurred: {e}")

        try:
            return False, e.completion, str(e)
        except Exception:
            return False, None, str(e)


if __name__ == "__main__":
    history = ImageChatHistory()
    history.add_system_msg("You are a helpful assistant.")
    history.add_user_msg("Hello, how are you?")

    success, content, error = call_model(history)
    if success:
        history.add_assistant_msg(content)

    else:
        print(error)

    print(history)
