
from openai import OpenAI

import textarena as ta 
from textarena.core import Agent

'''
class Agent(ABC):
    """ Generic agent class that defines the basic structure of an agent """
    @abstractmethod
    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.
        """
        pass
'''

class myAgent(Agent):
    def __init__(self, base_url: str, api_key: str):
        self.client = OpenAI(
            api_key=api_key, 
            base_url=base_url
        )

    def __call__(self, observation: str) -> str:
        response = self.client.chat.completions.create(
            model="Qwen3-4B",
            messages=[{"role": "user", "content": observation}],
        )
        print(response.choices[0].message)
        return response.choices[0].message.content.strip()

agent = myAgent(base_url="http://localhost:7024/v1", api_key="sk-KpWVukNiRPy2QTYA6fgH3CUwN53Twf")

agent("Hello, how are you?")

agents = {
    0: ta.agents.HumanAgent(),
    1: ta.agents.HumanAgent(),
    # 1: agent
    # 1: ta.agents.OpenRouterAgent(model_name="google/gemini-2.0-flash-001"),
}

# initialize the environment
env = ta.make(env_id="TicTacToe-v0")

env.reset(num_players=len(agents))

# main game loop
done = False 

while not done:
  player_id, observation = env.get_observation()
  action = agents[player_id](observation)
  done, step_info = env.step(action=action)

rewards, game_info = env.close()
print(rewards)
print(game_info)
