import json
import os
import openai
from typing import Optional
from openai import OpenAI
import google.generativeai as genai
import torch
import anthropic
from anthropic import Anthropic


# Function to create a prompt for OpenAI to check if the guess is correct
def prompt_to_openai_for_guess_correct(game_type, guesser_response, entity_name):

    """
    Creates a prompt for OpenAI to check if the guesser's response is correct.
    Args:
        game_type (str): The type of game, either "celebrities" or "things".
        guesser_response (str): The guesser's response to evaluate.
        entity_name (str): The name of the entity being guessed.
    Returns:
        str: The formatted OpenAI prompt.
    """


    if game_type == "celebrities":
        openai_prompt = f"""
        You are evaluating a guess made by a player in a 20-questions game. The player is attempting to guess the name of a celebrity.

        Here is the guesser's response: "{guesser_response}"

        Your task is to check if the guesser's response includes the **correct name of the celebrity** they are trying to guess. The name of the celebrity in English is: "{entity_name}". 

        The guess can be in any language or transliterated. It does not need to match the full name exactly (e.g., middle names can be neglected), but it must directly reference the celebrity's name. The guess must **explicitly include the celebrity's name** and cannot just describe their characteristics, actions, or attributes.

        Respond strictly in the following format:
        Response: "Yes"/"No"
        """
    
    elif game_type == "things":
        openai_prompt = f"""
        You are evaluating a guess made by a player in a 20-questions game. The player is attempting to guess the name of an object or thing.

        Here is the guesser's response: "{guesser_response}"

        Your task is to check if the guesser's response includes the **exact name of the object or thing** they are trying to guess. The name of the object or thing in English is: "{entity_name}". 

        The guess can be in any language or transliterated. Ensure the guess explicitly includes the correct name of the object or thing and not just descriptions, attributes, or vague references.

        Respond strictly in the following format:
        Response: "Yes"/"No"
        """

    return openai_prompt

# Function to create a prompt for OpenAI to check if the guesser has given up
def prompt_to_openai_for_giveup_check(game_type, guesser_response):

    """
    Creates a prompt for OpenAI to check if the guesser's response indicates they have given up.
    Args:
        game_type (str): The type of game, either "celebrities" or "things".
        guesser_response (str): The guesser's response to evaluate.
    Returns:
        str: The formatted OpenAI prompt.
    """


    if game_type == "celebrities":
        openai_prompt = f"""
        You are analyzing a player's response in a 20-questions guessing game. The player is attempting to guess the name of a **celebrity**.

        **Player's Response:** "{guesser_response}"

        Task:
        Determine whether the player's response **explicitly or implicitly** indicates that they have given up (i.e., they admit they are unable to guess the celebrity).  

        Important Notes:
        - The response **can be in any language** 

        A response should be considered as "giving up" if it:
        - Clearly states **inability to make a guess** and doesn't proceed to guess or ask question in the same response (e.g., "I don’t know," "I give up," "I can't figure it out," "No idea").  
        - Asks the other party to **reveal the answer** (e.g., "Tell me the answer").  
        - Uses words or phrases that strongly imply defeat or frustration and doesn't proceed to guess or ask question in the same response (e.g., "I surrender," "I have no guess," "I quit").  
        - Doesn't make a guess and asks a question/responds with something not related to the game (e.g., "What's the weather like?"). 

        A response should NOT be considered as "giving up" if it:
        - Makes an actual guess and asks if it is correct (e.g., "Is it Brad Pitt?").  
        - Asks a clarifying question to narrow down the options (e.g., "Is the person alive?").
        - Expresses uncertainty but still attempts to reason (e.g., "Not sure, but is it a musician?").
        - Clearly states **inability to make a guess** but proceeds to guess or ask question in the same response ("I can't figure it out, but is it a singer?").
        - Uses words or phrases that strongly imply defeat or frustration but proceeds to guess or ask question in the same response (e.g. "I have no guess, but is it Lionel Messi?").

        Answer Yes if the player's response indicates that they have given up, and No if it does not.

        Strictly respond in the following format:
        Response: "Yes"/"No"
        """

    elif game_type == "things":
        openai_prompt = f"""
        You are analyzing a player's response in a 20-questions guessing game. The player is attempting to guess the name of an **object or thing**.

        **Player's Response:** "{guesser_response}"

        Task:
        Determine whether the player's response **explicitly or implicitly** indicates that they have given up (i.e., they admit they are unable to guess the object or thing).

        Important Note:
        - The response **can be in any language**
        
        A response should be considered as "giving up" if it:
        - Clearly states **inability to make a guess** and doesn't proceed to guess or ask question in the same response (e.g., "I don’t know," "I give up," "I can't figure it out," "No idea").
        - Asks the other party to **reveal the answer** (e.g., "Tell me the answer").
        - Uses words or phrases that strongly imply defeat or frustration and doesn't proceed to guess or ask question in the same response (e.g., "I surrender," "I have no guess," "I quit").
        - Doesn't make a guess and asks a question/responds with something not related to the game (e.g., "What's the weather like?").

       A response should NOT be considered as "giving up" if it:
        - Makes an actual guess and asks if it is correct (e.g., "Is it a chair?").
        - Asks a clarifying question to narrow down the options (e.g., "Is it a living thing?").
        - Expresses uncertainty but still attempts to reason (e.g., "Not sure, but is it a living thing?").
        - Clearly states **inability to make a guess** but proceeds to guess or ask question in the same response ("I can't figure it out, but is it a living thing?").
        - Uses words or phrases that strongly imply defeat or frustration but proceeds to guess or ask question in the same response (e.g. "I have no guess, but is it a chair?").

        Answer Yes if the player's response indicates that they have given up, and No if it does not.

        Strictly respond in the following format:
        Response: "Yes"/"No"
        """

    return openai_prompt

class Q20Game:
    def __init__(
        self,
        item: str,
        answerer_model: str = "gpt-4o-mini",
        guesser_model: str = "gpt-4o-mini",
        num_turns: int = 20,
        temperature: float = 0.8,
        guesser_tokenizer=None,
        openai_api: bool = True,
        openai_api_key: Optional[str] = None,
        guesser_kargs={},
        language = "english",
        game_type = "things"
    ) -> None:
        self.item = item
        self.answerer_model = answerer_model
        self.guesser_model = guesser_model
        self.num_turns = num_turns
        self.temperature = temperature
        self.openai_api = openai_api
        self.guesser_tokenizer = guesser_tokenizer
        self.guesser_kargs = guesser_kargs
        self.language = language
        self.game_type = game_type
        

        if isinstance(guesser_model, str) and "gemini" in guesser_model:
            self.client_gemini = genai.Client(api_key=os.environ["GENAI_API_KEY"])

        # Read the prompts.json file
        with open("prompts.json") as f:
            prompts = json.load(f)
        
        self.prompts = prompts[self.game_type][self.language]
        
        self.vicuna_prompt = self.prompts['initial_system_prompt']
        self.first_user_utterance = self.prompts['guesser_prompt']
        self.guesser_win = False

        self.client_openai = OpenAI(
            api_key=os.environ.get("OPENAI_API_KEY"),
        )

        if isinstance(guesser_model, str) and "claude" in guesser_model:
            self.client_anthropic = Anthropic()

        self.guesser_messages = []

    # Function to call OpenAI models
    def api_call_openai(self, model_name, prompt, temp, max_tokens, messages=None):

        """
        Calls the OpenAI API to get a response based on the model name, prompt, temperature, and max tokens.
        If messages are provided, it uses them instead of the prompt.
        Args:
            model_name (str): The name of the OpenAI model to use.
            prompt (str): The prompt to send to the model.
            temp (float): The temperature for the model's response.
            max_tokens (int): The maximum number of tokens for the response.
            messages (Optional[List[dict]]): A list of messages to use instead of the prompt.
        Returns:
            str: The content of the response from the model, stripped of leading/trailing whitespace and asterisks.
        """

        if messages is None:
            # Create messages list using prompt
            messages = [
                {
                    "role": "user",
                    "content": f"{prompt}",
                }
            ]
        
        # Check if the model name starts with "gpt" or "o" to determine the API call type
        if model_name.startswith("gpt"):
            try:
                response = self.client_openai.chat.completions.create(
                model= model_name,
                temperature = temp,
                messages=messages,
                max_tokens= max_tokens
            )
                return response.choices[0].message.to_dict()["content"].strip().strip("**").strip("*").strip()
                
            except Exception as e:
                return None

        # Check if the model name starts with "o" for o series models
        elif model_name.startswith("o"):
            try:
                response = self.client_openai.chat.completions.create(
                model= model_name,
                temperature = temp,
                messages=messages,
                max_completion_tokens= max_tokens
            )
                return response.choices[0].message.to_dict()["content"].strip().strip("**").strip("*").strip()
                
            except Exception as e:
                return None
    
    # Function for API call to gemini
    def api_call_gemini(self, model_name, prompt, temp, max_tokens):

        """
        Calls the Gemini API to get a response based on the model name, prompt, temperature, and max tokens.
        Args:
            model_name (str): The name of the Gemini model to use.
            prompt (str): The prompt to send to the model.
            temp (float): The temperature for the model's response.
            max_tokens (int): The maximum number of tokens for the response.
        Returns:
            str: The content of the response from the model, stripped of leading/trailing whitespace.
        """

        # If the model name contains "2.5", it's a reasoning model and has a different max token limit
        if "2.5" in model_name:
            max_tokens = 10000

        try:
            response = self.client_gemini.models.generate_content(
                model= model_name,
                contents=[prompt],
                config=genai.types.GenerateContentConfig(
                    max_output_tokens=max_tokens,
                    temperature=temp
                )
            )

            # print("Response Text: ", response.text)
            return response.text
        except Exception as e:
            print(e)
            return None
    
    # Function for API call to Claude models
    def api_call_anthropic(self, model_name, messages, temp, max_tokens):
        
        """
        Calls the Anthropic API to get a response based on the model name, messages, temperature, and max tokens.
        Args:
            model_name (str): The name of the Anthropic model to use.
            messages (List[dict]): A list of messages to send to the model.
            temp (float): The temperature for the model's response.
            max_tokens (int): The maximum number of tokens for the response.
        Returns:
            str: The content of the response from the model, stripped of leading/trailing whitespace and asterisks.
        """



        # Convert to Claude-compatible format
        claude_messages = []
        for msg in messages:
            claude_messages.append({
                "role": msg["role"],
                "content": [
                    {
                        "type": "text",
                        "text": msg["content"]
                    }
                ]
            })
        try:
            # API call
            response = self.client_anthropic.messages.create(
                model= model_name,
                max_tokens=max_tokens,
                temperature=temp,
                messages = claude_messages,
            )

            return response.content[0].text.strip().strip("**").strip("*").strip()
        
        except Exception as e:
            print(e)
            return None


    # Function to call the guesser model
    def guesser(self, messages):

        """
        Calls the guesser model to get a response
        Args:
            messages (List[dict]): A list of messages to send to the guesser model.
        Returns:
            dict: A dictionary containing the content of the guesser's response.
        """

        # print("Guesser model: ", self.guesser_model)

        # Check if the guesser model is a claude model
        if isinstance(self.guesser_model, str) and self.guesser_model.startswith(
            "claude"
        ):
            
            times_runned = 0
            while(times_runned<5):

                content = self.api_call_anthropic(
                    self.guesser_model,
                    messages,
                    self.temperature,
                    100,
                )

                # Check if response is None or empty string
                if content is None or content == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        content = "N/A"
                        print("Response from Anthropic is None! Breaking the loop")
                        break
                    continue
                break

            content = content.replace("\n", ". ")

            return {
                "role": "assistant",
                "content": content,
            }


        # Check if the guesser model is a huggingface model
        elif not isinstance(self.guesser_model, str):
            # """Wraps hf's `generate` adding some specific method's defaults"""
            prompt = self.dialog_history() + " ASSISTANT:"
            # Ensure tokenizer has a pad token
            if self.guesser_tokenizer.pad_token is None:
                self.guesser_tokenizer.pad_token = self.guesser_tokenizer.eos_token

            # Encode with attention mask
            encoded_input = self.guesser_tokenizer(
                prompt, return_tensors="pt", padding=True
            )

            input_ids = encoded_input["input_ids"].to(self.guesser_model.device)
            attention_mask = encoded_input["attention_mask"].to(self.guesser_model.device)

            with torch.no_grad():
                gen = self.guesser_model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **self.guesser_kargs,
                )
                gen_str = (
                    self.guesser_tokenizer.decode(gen[0][input_ids[0].shape[0] :])
                    .split("</s>")[0]
                    .split("USER")[0]
                    .split("<|eot_id|>")[0]
                    .lstrip()
                    .strip() 
                )

            # Check if there is a </s> in the gen_str, t there is then split it and take the previous part
            if "</s>" in gen_str:
                gen_str = gen_str.split("</s>")[0].strip()

            # Check if there are any /n, if yes then replace it with ". "
            gen_str = gen_str.replace("\n", ". ")

            gen_str = gen_str.replace("ASSISTANT", "")


            return {
                "role": "assistant",
                "content": gen_str,
            }
        
        # Check if the guesser model is a gemini model
        elif (isinstance(self.guesser_model, str)) and ("gemini" in self.guesser_model):
                                                    
            prompt = ""
            for item in messages:
                if item["role"].upper() == "USER":
                    prompt += f"USER: {item['content']}\n"
                elif item["role"].upper() == "ASSISTANT":
                    prompt += f"ASSISTANT: {item['content']}\n"

            prompt+= "ASSISTANT:"

            times_runned = 0
            while(times_runned<5):
                response = self.api_call_gemini(self.guesser_model, prompt, self.temperature, 100)

                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from Gemini is None! Breaking the loop")
                        break
                    continue
                break
            
            # Split the response by first USER or User and then take the first part
            if "USER" in response:
                response = response.split("USER")[0].strip()
                # Remove any \n
                response = response.replace("\n", ". ")
            
            if "User" in response:
                response = response.split("User")[0].strip()
                # Remove any \n
                response = response.replace("\n", ". ")

            response = response.replace("\n", ". ")

            # Remove if there is any ASSISTANT string in response
            response = response.replace("ASSISTANT", "")
            
            return {
                "role": "assistant",
                "content": response.strip(),
            }

        # Check if the guesser model is a GPT model
        elif isinstance(self.guesser_model, str) and "gpt" in self.guesser_model:

            times_runned = 0
            while(times_runned<5):

                content = self.api_call_openai(
                    self.guesser_model,
                    "",
                    self.temperature,
                    100,
                    messages=messages,
                )

                # Check if response is None or empty string
                if content is None or content == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        content = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break

            content = content.replace("\n", ". ")

            return {
                "role": "assistant",
                "content": content,
            }

        # Check if the guesser model is an OpenAI o series model
        elif isinstance(self.guesser_model, str) and self.guesser_model.startswith("o"):
            times_runned = 0
            while(times_runned<5):
                content = self.api_call_openai(
                self.guesser_model,
                "",
                1.0,
                10000,
                messages=messages,
                )

                content = content.replace("\n", ". ")

                # Check if response is None or empty string
                if content is None or content == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        content = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break


            return {
                "role": "assistant",
                "content": content,
            }


    # Function to get the dialog history in the format required by the guesser model
    def dialog_history(self):

        """
        Returns the dialog history in the format required by the guesser model.
        Returns:
            str: The dialog history as a string.
        """


        history = self.vicuna_prompt + " "
        for item in self.guesser_messages:
            if item["role"].upper() == "USER":
                history += "USER: " + item["content"]
            elif item["role"].upper() == "ASSISTANT":
                history += " " + "ASSISTANT: " + item["content"] + "</s>"
        return history

    # Function in which the whole game is played
    def game_play(self, path, user_mode=False):

        """
        Plays the 20-questions game with the guesser and answerer.
        Args:
            path (str): The path where the game session will be saved.
            user_mode (bool): If True, the user can input their own questions. If False, the model will ask questions.
        Returns:
            bool: True if the guesser wins, False if the guesser loses or gives up.
        """

        # Check if the game is already played by looking if the directory exists
        if os.path.exists(os.path.join(path, f"{self.item}.txt")):
            print("Game already played for ", self.item)
            return
    
        self.reset()
        # print(f"Item: {self.item}")
        for t in range(self.num_turns):
            # System asking a question
            if (not user_mode) or user_mode is None:
                guesser_msg = self.guesser(self.guesser_messages)
        
            else:
                user_q = input(
                    f"Type in your questions for turn {t+1}. (e.g. Is it a living thing?)\n"
                )
                guesser_msg = {"role": "assistant", "content": user_q}
            self.guesser_messages.append(guesser_msg)
            guesser_question = guesser_msg["content"].strip()

            if t == self.num_turns - 1:
                self.guesser_messages[-1]["content"] = (
                    self.guesser_messages[-1]["content"] + self.prompts['guesser_asking_if_guess_is_correct']
                )

            usr_msg = self.answerer(guesser_question)


            self.guesser_messages.append(
                {"role": "user", "content": f"{usr_msg['content'].strip()}"}
            )

            if "bingo" in usr_msg["content"].lower():
                self.guesser_win = True
                return True
            
            # Check if given up
            if "you have given up" in usr_msg["content"].lower():
                self.guesser_win = False
                return False
            
            if t == self.num_turns - 2:
                self.guesser_messages[-1]["content"] = (
                    self.guesser_messages[-1]["content"]
                    + self.prompts['telling_guesser_it_has_to_guess_now']
                )

        return False

    # Save the game in a text file in the given path
    def save_session(self, path):

        """
        Saves the game session to a text file in the specified path.
        Args:
            path (str): The path where the game session will be saved.
        Returns:
            dict: A dictionary containing the item name as the key and a dictionary with "guesser" and "answerer" as keys and their respective messages/responses as values.
        """


        # Print the conversation
        if not os.path.exists(path):
            os.makedirs(path)
        output_file = os.path.join(path, f"{self.item}.txt")
        # Check if the game is already played by looking if the directory exists
        if os.path.exists(output_file):
            print(f"Game already played for {self.item}")
            # Reed the file and store the content in a dictionary with the item as key and value as another dictionary with the "guesser" and "answerer" as keys and values as the list of messages/responses
            with open(output_file, "r") as out_f:
                lines = out_f.readlines()
                
            data = {}
            # Get the item name
            item_name = self.item

            guesser_msgs = []
            answerer_msgs = []

            for line in lines[2:]:
                line = line.strip()
                if not line:
                    continue
                if "Assistant:" in line:
                    # Extract assistant message (question)
                    guesser_msgs.append(line.split("Assistant: ", 1)[1])
                elif "User:" in line:
                    # Extract user message (answer)
                    answerer_msgs.append(line.split("User: ", 1)[1])
            data[item_name] = {
                "guesser": guesser_msgs,
                "answerer": answerer_msgs
            }
            return data
        

        with open(output_file, "w") as out_f:
            out_f.write(f"item: {self.item}\n")
            for t, message in enumerate(self.guesser_messages):
                out_f.write(
                    f"Turn {(t+1)//2}, {message['role'].capitalize()}: {message['content'].lstrip()}\n"
                )
        
        # Read the file and store the content in a dictionary with the item as key and value as another dictionary with the "guesser" and "answerer" as keys and values as the list of messages
        # Open the file and read the content
        with open(output_file, "r") as out_f:
            lines = out_f.readlines()

        data = {}
        # Get the item name
        item_name = self.item

        guesser_msgs = []
        answerer_msgs = []

        for line in lines[2:]:  # Skip the first two lines
            line = line.strip()
            if not line:
                continue
            if "Assistant:" in line:
                # Extract assistant message (question)
                guesser_msgs.append(line.split("Assistant: ", 1)[1])
            elif "User:" in line:
                # Extract user message (answer)
                answerer_msgs.append(line.split("User: ", 1)[1])

        data[item_name] = {
            "guesser": guesser_msgs,
            "answerer": answerer_msgs
        }

        return data


    # Function to call the answerer model
    def answerer(self, question):

        """
        Calls the answerer model to get a response based on the question asked by the guesser.
        Args:
            question (str): The question asked by the guesser.
        Returns:
            dict: A dictionary containing the content of the answerer's response.
        """
    


        user_messages = [
            {
                "role": "user",
                "content": self.prompts['answerer_system_prompt'].format(item = self.item),
            },
            {
                "role": "user",
                "content": self.prompts['answerer_prompt'].format(question = question, item = self.item),
            },
        ]


        times_runned = 0


        # Check if the guess given my the guesser is correct or not
        while(times_runned<5):

            response_from_openai_guess_check = self.api_call_openai("gpt-4o-mini",  prompt_to_openai_for_guess_correct(self.game_type, question, self.item), 0.2, 30)
            if response_from_openai_guess_check is None:
                print("Response is None! Retrying")
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_guess_check = "N/A"
                    print("Response from OpenAI is None! Breaking the loop")
                    break
                continue
                
            if response_from_openai_guess_check.startswith("Response:") or response_from_openai_guess_check.startswith("***Response:***") or response_from_openai_guess_check.startswith("**Response:**"):
                response_from_openai_guess_check = response_from_openai_guess_check.split(":")[1].strip().lower().strip('"').strip()
                # Check if it is omly yes or no
                if response_from_openai_guess_check == "yes" or response_from_openai_guess_check == "no":
                    break
                else:
                    times_runned += 1
                    if (times_runned == 5):
                        response_from_openai_guess_check = "N/A"
                        print("Response from OpenAI is not containing only yes or no! Breaking the loop")
                        break
                    continue

            if not response_from_openai_guess_check.startswith("Response:") or not response_from_openai_guess_check.startswith("***Response:***") or not response_from_openai_guess_check.startswith("**Response:**"):
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_guess_check = "N/A"
                    print("Response from OpenAI is not in starting with response! Breaking the loop")
                    break
                continue
        
        # Giveup check
        times_runned = 0

        while(times_runned<5):

            response_from_openai_check_giveup = self.api_call_openai("gpt-4o-mini",  prompt_to_openai_for_giveup_check(self.game_type, question), 0.2, 30)
            if response_from_openai_check_giveup is None:
                print("Response is None! Retrying")
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_check_giveup = "N/A"
                    print("Response from OpenAI is None! Breaking the loop")
                    break
                continue
                
            if response_from_openai_check_giveup.startswith("Response:") or response_from_openai_check_giveup.startswith("***Response:***") or response_from_openai_check_giveup.startswith("**Response:**"):
                response_from_openai_check_giveup = response_from_openai_check_giveup.split(":")[1].strip().lower().strip('"').strip()
                # Check if it is omly yes or no
                if response_from_openai_check_giveup == "yes" or response_from_openai_check_giveup == "no":
                    break
                else:
                    times_runned += 1
                    if (times_runned == 5):
                        response_from_openai_check_giveup = "N/A"
                        print("Response from OpenAI is not containing only yes or no! Breaking the loop")
                        break
                    continue

            if not response_from_openai_check_giveup.startswith("Response:") or not response_from_openai_check_giveup.startswith("***Response:***") or not response_from_openai_check_giveup.startswith("**Response:**"):
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_check_giveup = "N/A"
                    print("Response from OpenAI is not in starting with response! Breaking the loop")
                    break
                continue

        # Check if the response_from_openai_check_giveup is yes
        if response_from_openai_check_giveup == "yes":
            return {"role": "user", "content": "You have given up!"}
        
        # Check if the answerer model is a GPT model
        if isinstance(self.guesser_model, str) and "gpt" in self.guesser_model:
            
            times_runned = 0
            while(times_runned<5):

                response = self.api_call_openai(
                    self.guesser_model,
                    "",
                    0.2,
                    30,
                    messages=user_messages,
                )
                # Check if response is None or empty string

                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            # return response.choices[0].message.to_dict()
            return {
                "role": "user",
                "content": response,
            }
        

        # Check if the guesser model is an OpenAI o series model
        elif isinstance(self.guesser_model, str) and self.guesser_model.startswith("o"):
            
            times_runned = 0
            while(times_runned<5):

                response = self.api_call_openai(
                    self.guesser_model,
                    "",
                    1.0,
                    10000,
                    messages=user_messages,
                )

                # Check if response is None or empty string
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break
            
            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            return {
                "role": "user",
                "content": response,
            }

        # Check if the guesser model is a claude model
        elif isinstance(self.guesser_model, str) and "claude" in self.guesser_model:

            times_runned = 0
            while(times_runned<5):

                response = self.api_call_anthropic(
                    self.guesser_model,
                    user_messages,
                    0.2,
                    30,
                )

                # Check if response is None or empty string
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}
            

            # return response.choices[0].message.to_dict()
            return {
                "role": "user",
                "content": response,
            }
        
        # Check if the guesser model is a gemini model
        elif isinstance(self.guesser_model, str) and "gemini" in self.guesser_model:

            prompt = ""
            for item in user_messages:
                if item["role"].upper() == "USER":
                    prompt += f"USER: {item['content']}\n"
                elif item["role"].upper() == "ASSISTANT":
                    prompt += f"ASSISTANT: {item['content']}\n"

            prompt+= "ASSISTANT:"

            times_runned = 0
            while(times_runned<5):

                response = self.api_call_gemini(self.guesser_model, prompt, 0.2, 30)
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from Gemini is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            return {"role": "user", "content": response.strip()}
        
        # Check if the guesser model is a huggingface model
        elif not isinstance(self.guesser_model, str):

            prompt = ""
            for item in user_messages:
                if item["role"].upper() == "USER":
                    prompt += f"USER: {item['content']}\n"
                elif item["role"].upper() == "ASSISTANT":
                    prompt += f"ASSISTANT: {item['content']}\n"

            prompt+= "ASSISTANT:"

            # Ensure tokenizer has a pad token
            if self.guesser_tokenizer.pad_token is None:
                self.guesser_tokenizer.pad_token = self.guesser_tokenizer.eos_token

            # Encode with attention mask
            encoded_input = self.guesser_tokenizer(
                prompt, return_tensors="pt", padding=True
            )

            input_ids = encoded_input["input_ids"].to(self.guesser_model.device)
            attention_mask = encoded_input["attention_mask"].to(self.guesser_model.device)

            with torch.no_grad():
                gen = self.guesser_model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **self.guesser_kargs,
                )
                gen_str = (
                    self.guesser_tokenizer.decode(gen[0][input_ids[0].shape[0] :])
                    .split("</s>")[0]
                    .split("USER")[0]
                    .split("<|eot_id|>")[0]
                    .lstrip()
                    .strip() 
                )

            # if "llama" in self.guesser_model:
            #     gen_str = self.api_call_llama(prompt, 0.2, 30)
            
            # elif "mistral" in self.guesser_model:
            #     gen_str = self.api_call_mistral(prompt, 30, 0.2)

            # Check if there is a </s> in the gen_str, t there is then split it and take the previous part
            if "</s>" in gen_str:
                gen_str = gen_str.split("</s>")[0].strip()

            # Check if there are any /n, if yes then replace it with ". "
            gen_str = gen_str.replace("\n", ". ")

            gen_str = gen_str.replace("ASSISTANT", "")

            gen_str = gen_str.strip()

            if response_from_openai_guess_check == "yes" and "yes" in gen_str.lower():
                return {"role": "user", "content": "Bingo!"}

            return {
                "role": "user",
                "content": gen_str,
            }

    # Function to reset the game
    def reset(self):
        
        """
        Resets the game by clearing the guesser messages and initializing the first user utterance.
        """

        # Initialize the conversation
        self.guesser_messages = [
            {
                "role": "user",
                "content": self.first_user_utterance,
            }
        ]


class Q20GameCelebrity(Q20Game):
    def __init__(self, item: str, background: str, language, game_type, **kwargs) -> None:
        super().__init__(item, **kwargs, language=language, game_type=game_type)
        self.background = background
        self.first_user_utterance = self.prompts['guesser_prompt']


    # Function to call the answerer model
    def answerer(self, question):

        """
        Calls the answerer model to get a response based on the question asked by the guesser.
        Args:
            question (str): The question asked by the guesser.
        Returns:
            dict: A dictionary containing the content of the answerer's response.
        """

        # if isinstance(self.guesser_model, str) and "gpt" in self.guesser_model:
        #     openai.api_base = "https://api.openai.com/v1"

        if len(self.background) < 10:
            user_messages = [
                {
                    "role": "user",
                    "content": self.prompts['answerer_system_prompt_without_background_of_celebrity_given'].format(item = self.item),
                },
                {
                    "role": "user",
                    "content": self.prompts['answerer_prompt_without_background_of_celebrity_given'].format(item = self.item, question = question),
                },
            ]
        else:
            user_messages = [
                {
                    "role": "user",
                    "content": self.prompts['answerer_system_prompt_with_background_of_celebrity_given'].format(item = self.item, background = self.background),
                },
                {
                    "role": "user",
                    "content": self.prompts['answerer_prompt_with_background_of_celebrity_given'].format(item = self.item, question = question),
                },
            ]

        times_runned = 0

        # Check if the guess given my the guesser is correct or not
        while(times_runned<5):

            response_from_openai_guess_check = self.api_call_openai("gpt-4o-mini",  prompt_to_openai_for_guess_correct(self.game_type, question, self.item), 0.2, 30)
            if response_from_openai_guess_check is None:
                print("Response is None! Retrying")
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_guess_check = "N/A"
                    print("Response from OpenAI is None! Breaking the loop")
                    break
                continue
            
            if response_from_openai_guess_check.startswith("Response:") or response_from_openai_guess_check.startswith("***Response:***") or response_from_openai_guess_check.startswith("**Response:**"):
                response_from_openai_guess_check = response_from_openai_guess_check.split(":")[1].strip().lower().strip('"').strip()
                # Check if it is omly yes or no
                if response_from_openai_guess_check == "yes" or response_from_openai_guess_check == "no":
                    break
                else:
                    times_runned += 1
                    print(f"Response from OpenAI is not containing only yes or no! {response_from_openai_guess_check}")
                    if (times_runned == 5):
                        response_from_openai_guess_check = "N/A"
                        print("Response from OpenAI is not containing only yes or no! Breaking the loop")
                        break
                    continue
            
            if not response_from_openai_guess_check.startswith("Response:") or not response_from_openai_guess_check.startswith("***Response:***") or not response_from_openai_guess_check.startswith("**Response:**"):
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_guess_check = "N/A"
                    print("Response from OpenAI is not in starting with response! Breaking the loop")
                    break
                continue

        # Giveup check
        times_runned = 0

        while(times_runned<5):
            
            response_from_openai_check_giveup = self.api_call_openai("gpt-4o-mini",  prompt_to_openai_for_giveup_check(self.game_type, question), 0.2, 30)
            if response_from_openai_check_giveup is None:
                print("Response is None! Retrying")
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_check_giveup = "N/A"
                    print("Response from OpenAI is None! Breaking the loop")
                    break
                continue
                
            if response_from_openai_check_giveup.startswith("Response:") or response_from_openai_check_giveup.startswith("***Response:***") or response_from_openai_check_giveup.startswith("**Response:**"):
                response_from_openai_check_giveup = response_from_openai_check_giveup.split(":")[1].strip().lower().strip('"').strip()
                # Check if it is omly yes or no
                if response_from_openai_check_giveup == "yes" or response_from_openai_check_giveup == "no":
                    break
                else:
                    times_runned += 1
                    if (times_runned == 5):
                        response_from_openai_check_giveup = "N/A"
                        print("Response from OpenAI is not containing only yes or no! Breaking the loop")
                        break
                    continue

            if not response_from_openai_check_giveup.startswith("Response:") or not response_from_openai_check_giveup.startswith("***Response:***") or not response_from_openai_check_giveup.startswith("**Response:**"):
                times_runned += 1
                if (times_runned == 5):
                    response_from_openai_check_giveup = "N/A"
                    print("Response from OpenAI is not in starting with response! Breaking the loop")
                    break
                continue


        # Check if the response_from_openai_check_giveup is yes
        if response_from_openai_check_giveup == "yes":
            return {"role": "user", "content": "You have given up!"}
            
        # Check if the answerer model is a GPT model
        if isinstance(self.guesser_model, str) and "gpt" in self.guesser_model:
            
            times_runned = 0
            while(times_runned<5):
                response = self.api_call_openai(
                    self.guesser_model,
                    "",
                    0.2,
                    30,
                    messages=user_messages,
                )

                # Check if response is None or empty string
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            # return response.choices[0].message.to_dict()
            return {
                "role": "user",
                "content": response,
            }
        
        # Check if the guesser model is an OpenAI o series model
        elif isinstance(self.guesser_model, str) and self.guesser_model.startswith("o"):
            
            times_runned = 0
            while(times_runned<5):

                response = self.api_call_openai(
                    self.guesser_model,
                    "",
                    1.0,
                    10000,
                    messages=user_messages,
                )

                # Check if response is None or empty string
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            return {
                "role": "user",
                "content": response,
            }

        # Check if the guesser model is a claude model
        elif isinstance(self.guesser_model, str) and "claude" in self.guesser_model:

            times_runned = 0

            while(times_runned<5):

                response = self.api_call_anthropic(
                    self.guesser_model,
                    user_messages,
                    0.2,
                    30,
                )

                # Check if response is None or empty string
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from OpenAI is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            # return response.choices[0].message.to_dict()
            return {
                "role": "user",
                "content": response,
            }
        
        # Check if the guesser model is a gemini model
        elif isinstance(self.guesser_model, str) and "gemini" in self.guesser_model:

            prompt = ""
            for item in user_messages:
                if item["role"].upper() == "USER":
                    prompt += f"USER: {item['content']}\n"
                elif item["role"].upper() == "ASSISTANT":
                    prompt += f"ASSISTANT: {item['content']}\n"

            prompt+= "ASSISTANT:"

            times_runned = 0
            while(times_runned<5):

                response = self.api_call_gemini(self.guesser_model, prompt, 0.2, 30)
                if response is None or response == "":
                    print("Response is None! Retrying")
                    times_runned += 1
                    if (times_runned == 5):
                        response = "N/A"
                        print("Response from Gemini is None! Breaking the loop")
                        break
                    continue
                break

            response = response.strip()

            if response_from_openai_guess_check == "yes" and "yes" in response.lower():
                return {"role": "user", "content": "Bingo!"}

            return {"role": "user", "content": response.strip()}
        
        # Check if the guesser model is a huggingface model
        elif not isinstance(self.guesser_model, str):

            prompt = ""
            for item in user_messages:
                if item["role"].upper() == "USER":
                    prompt += f"USER: {item['content']}\n"
                elif item["role"].upper() == "ASSISTANT":
                    prompt += f"ASSISTANT: {item['content']}\n"

            prompt+= "ASSISTANT:"

            # Ensure tokenizer has a pad token
            if self.guesser_tokenizer.pad_token is None:
                self.guesser_tokenizer.pad_token = self.guesser_tokenizer.eos_token

            # Encode with attention mask
            encoded_input = self.guesser_tokenizer(
                prompt, return_tensors="pt", padding=True
            )

            input_ids = encoded_input["input_ids"].to(self.guesser_model.device)
            attention_mask = encoded_input["attention_mask"].to(self.guesser_model.device)

            with torch.no_grad():
                gen = self.guesser_model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **self.guesser_kargs,
                )
                gen_str = (
                    self.guesser_tokenizer.decode(gen[0][input_ids[0].shape[0] :])
                    .split("</s>")[0]
                    .split("USER")[0]
                    .split("<|eot_id|>")[0]
                    .lstrip()
                    .strip() 
                )


            # Check if there is a </s> in the gen_str, t there is then split it and take the previous part
            if "</s>" in gen_str:
                gen_str = gen_str.split("</s>")[0].strip()

            # Check if there are any /n, if yes then replace it with ". "
            gen_str = gen_str.replace("\n", ". ")

            gen_str = gen_str.replace("ASSISTANT", "")

            gen_str = gen_str.strip()

            if response_from_openai_guess_check == "yes" and "yes" in gen_str.lower():
                return {"role": "user", "content": "Bingo!"}

            return {
                "role": "user",
                "content": gen_str,
            }

    # Function to reset the game
    def reset(self):

        """
        Resets the game by clearing the guesser messages and initializing the first user utterance.
        """

        # Initialize the conversation
        self.guesser_messages = [
            {
                "role": "user",
                "content": self.first_user_utterance,
            }
        ]