from abc import ABC, abstractmethod
import openai, os, time
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from typing import Optional 

from textarena.core import Agent
import textarena as ta 

__all__ = [
    "HumanAgent",
    "OpenRouterAgent",
    "HFLocalAgent",
    "CerebrasAgent"
]

    
class HumanAgent(Agent):
    """
    Human agent class that allows the user to input actions manually.
    """
    def __init__(self):
        """
        Initialize the human agent.
        
        Args:
            model_name (str): The name of the model.
        """
        super().__init__()

    def __call__(
        self, 
        observation: str
    ) -> str:
        """
        Process the observation and return the action.
        
        Args:
            observation (str): The input string to process.
            
        Returns:
            str: The response generated by the agent.
        """
        print(observation)
        return input("Please enter the action: ")


# class OpenRouterAgent(Agent):
#     """
#     GPT agent class that uses the OpenRouter API to generate responses.
#     """
#     def __init__(
#         self, 
#         model_name: str,
#         system_prompt: Optional[str]=None,
#         verbose: Optional[bool]=False
#     ):
#         """
#         Initialize the GPT agent.
        
#         Args:
#             model_name (str): The name of the model.
#             system_prompt (str): The system prompt to use (default: "You are a competitive game player.").
#         """
#         super().__init__()
#         self.model_name = model_name
#         self.verbose = verbose

#         ## Set the OpenAI API key
#         openai.api_key = os.getenv("OPENAI_API_KEY")
#         if not openai.api_key:
#             raise ValueError("OPENAI API key not found. Please set the OPENAI_API_KEY environment variable.")
        
#         ## Initialize the OpenAI client
#         self.client = openai.OpenAI(base_url="https://openrouter.ai/api/v1")

#         ## Set the system prompt
#         if system_prompt is None:
#             self.system_prompt = "You are a competitive game player. Make sure you read the game instructions carefully, and always follow the required format."
#         else:
#             self.system_prompt = system_prompt

    
#     def __call__(
#         self, 
#         observation: str
#     ) -> str:
#         """
#         Process the observation using the OpenAI model and return the action.
        
#         Args:
#             observation (str): The input string to process.
        
#         Returns:
#             str: The response generated by the model.
#         """
#         assert isinstance(observation, str), \
#             f"When uisng OpenRouter, the observation must be a string. Received type: {type(observation)}"
#         try:
#             response = self.client.chat.completions.create(
#                 model=self.model_name,
#                 messages=[
#                     {"role": "system", "content": self.system_prompt},
#                     {"role": "user", "content": observation}
#                 ],
#                 # max_tokens=150, ## optional
#                 n=1,
#                 stop=None,
#                 temperature=0.7,
#             )
#             if self.verbose:
#                 print(f"\nObservation: {observation}\n Response: {response}")
#             # Extract the assistant's reply
#             action = response.choices[0].message.content.strip()
#             return action
#         except Exception as e:
#             time.sleep(10_000)
#             # raise E
#             return f"An error occurred: {e}"



class OpenRouterAgent(Agent):
    """
    GPT agent class that uses the OpenRouter API to generate responses.
    """
    def __init__(
        self, 
        model_name: str,
        system_prompt: Optional[str]=None,
        verbose: Optional[bool]=False
    ):
        """
        Initialize the GPT agent.
        
        Args:
            model_name (str): The name of the model.
            system_prompt (str): The system prompt to use (default: "You are a competitive game player.").
        """
        super().__init__()
        self.model_name = model_name
        self.verbose = verbose

        ## Set the OpenAI API key
        openai.api_key = os.getenv("OPENAI_API_KEY")
        if not openai.api_key:
            raise ValueError("OPENAI API key not found. Please set the OPENAI_API_KEY environment variable.")
        
        ## Initialize the OpenAI client
        self.client = openai.OpenAI(base_url="https://openrouter.ai/api/v1")

        ## Set the system prompt
        if system_prompt is None:
            self.system_prompt = "You are a competitive game player. Make sure you read the game instructions carefully, and always follow the required format."
        else:
            self.system_prompt = system_prompt

    
    def __call__(
        self, 
        observation: str
    ) -> str:
        """
        Process the observation using the OpenAI model and return the action.
        
        Args:
            observation (str): The input string to process.
        
        Returns:
            str: The response generated by the model.
        """
        assert isinstance(observation, str), \
            f"When uisng OpenRouter, the observation must be a string. Received type: {type(observation)}"
        
        for _ in range(5):
            try:
                response = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=[
                        {"role": "system", "content": self.system_prompt},
                        {"role": "user", "content": observation}
                    ],
                    # max_tokens=150, ## optional
                    n=1,
                    stop=None,
                    # temperature=0.7,
                )
                if self.verbose:
                    print(f"\nObservation: {observation}\n Response: {response}")
                # Extract the assistant's reply
                action = response.choices[0].message.content.strip()
                return action
            except Exception as e:
                print(f"\n\n{e}")
                time.sleep(5)

            print(f"Model failed. going to sleep")
            time.sleep(10_000)
            # try:
            #     response = self.client.chat.completions.create(
            #         model=self.model_name,
            #         messages=[
            #             {"role": "system", "content": self.system_prompt},
            #             {"role": "user", "content": observation}
            #         ],
            #         # max_tokens=150, ## optional
            #         n=1,
            #         stop=None,
            #         temperature=0.7,
            #     )
            #     if self.verbose:
            #         print(f"\nObservation: {observation}\n Response: {response}")
            #     # Extract the assistant's reply
            #     action = response.choices[0].message.content.strip()
            #     return action
            # except Exception as e:
            #     print(f"\n\n{e}")
            #     time.sleep(10_000)
            #     # raise E
            #     return f"An error occurred: {e}"
            


class HFLocalAgent(Agent):
    """
    Hugging Face local agent class that uses the Hugging Face Transformers library.
    """
    def __init__(
        self, 
        model_name: str, 
        device: str = "auto",
        quantize: bool = False
    ):
        """
        Initialize the Hugging Face local agent.
        
        Args:
            model_name (str): The name of the model.
            quantize (bool): Whether to load the model in 8-bit quantized format (default: False).
        """
        super().__init__()
        ## Initialize the Hugging Face model and tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name, 
            )
        
        if quantize:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name, 
                load_in_8bit=True,
                device_map=device,
                )
        else:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                device_map=device,
                )

        ## Initialize the Hugging Face pipeline
        self.pipeline = pipeline(
            'text-generation',
            max_new_tokens=500,
            model=self.model, 
            tokenizer=self.tokenizer, 
            )
    
    def __call__(
        self, 
        observation: str
    ) -> str:
        """
        Process the observation using the Hugging Face model and return the action.
        
        Args:
            observation (str): The input string to process.
        
        Returns:
            str: The response generated by the model.
        """
        # Generate a response
        try:
            response = self.pipeline(
                observation, 
                num_return_sequences=1, 
                return_full_text=False,
            )
            # Extract and return the text output
            action = response[0]['generated_text'].strip()
            return action
        except Exception as e:
            return f"An error occurred: {e}"



class CerebrasAgent(Agent):
    """
    Cerebras agent class that uses the Cerebras API to generate responses.
    """

    def __init__(self, model_name: str, system_prompt: str | None = None):
        """
        Initialize the Cerebras agent.

        Args:
            model_name (str): The name of the model.
            system_prompt (str): The system prompt to use (default: "You are a competitive game player.").
        """
        super().__init__()
        self.model_name = model_name
        
        from cerebras.cloud.sdk import Cerebras
        self.client = Cerebras(
            # This is the default and can be omitted
            api_key=os.getenv("CEREBRAS_API_KEY"),
        )

        ## Set the system prompt
        if system_prompt is None:
            self.system_prompt = "You are a competitive game player. Make sure you read the game instructions carefully, and always follow the required format."
        else:
            self.system_prompt = system_prompt

    def __call__(self, observation: str) -> str:
        """
        Process the observation using the Cerebras model and return the action.

        Args:
            observation (str): The input string to process.

        Returns:
            str: The response generated by the model.
        """
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": observation},
                ],
                top_p=0.9,
                temperature=0.9,
            )
            # Extract the assistant's reply
            action = response.choices[0].message.content.strip()
            return action
        except Exception as e:
            return f"An error occurred: {e}"