import datetime
import os
import random
import traceback
from threading import Thread
from typing import List, Optional
from loguru import logger

from agent import KEN_GREEN, KEN_RED, TextRobot, VisionRobot, TextLocalRobot, SocketConfig
from agent.config import MODELS
from diambra.arena import (
    EnvironmentSettingsMultiAgent,
    RecordingSettings,
    SpaceTypes,
    make,
)
from rich import print

from agent.robot import Robot
import time

def generate_random_model(openai: bool = False, mistral: bool = True):
    models_available = []

    for model, models in MODELS.items():
        if openai and model == "OPENAI":
            models_available.extend(models)
        if mistral and model == "MISTRAL":
            models_available.extend(models)

    random.seed()
    # Generate a pair of random two models
    random_model = random.choice(models_available)

    return random_model


class Player:
    nickname: str
    model: str
    robot: Optional[Robot] = None
    temperature: float = 0.7

    def verify_provider_name(self):
        if self.model.startswith("openai"):
            assert (
                os.environ.get("OPENAI_API_KEY") is not None
            ), "OpenAI API key not set"
        if self.model.startswith("mistral"):
            assert (
                os.environ.get("MISTRAL_API_KEY") is not None
            ), "Mistral API key not set"
        if self.model.startswith("cerebras"):
            assert (
                os.environ.get("CEREBRAS_API_KEY") is not None
            ), "Cerebras API key not set"
        if self.model.startswith("remote"):
            return


class Player1(Player):
    def __init__(
        self,
        nickname: str,
        model: str,
        robot_type: str = "text",
        temperature: float = 0.7,
        host: str = "localhost",
        port: int = 38001,
        max_client: int = 1,
        serving_method: str = "remote",
    ):
        self.nickname = nickname
        self.model = model
        self.robot_type = robot_type
        self.temperature = temperature

        if robot_type == "vision":
            self.robot = VisionRobot(
                action_space=None,
                character="Ken",
                side=0,
                character_color=KEN_RED,
                model=self.model,
                ennemy_color=KEN_GREEN,
                only_punch=os.getenv("TEST_MODE", False),
                temperature=self.temperature,
                sleepy=False,
                player_nb=1,
            )
        elif robot_type == "text":
            self.robot = TextRobot(
                action_space=None,
                character="Ken",
                side=0,
                character_color=KEN_RED,
                ennemy_color=KEN_GREEN,
                only_punch=os.getenv("TEST_MODE", False),
                temperature=self.temperature,
                sleepy=False,
                model=self.model,
                player_nb=1,
            )
        elif robot_type == "remote":
            socket_config = SocketConfig(
                Host=host,
                Port=port,
                max_client=max_client,
            )

            self.robot = TextLocalRobot(
                action_space=None,
                character="Ken",
                side=0,
                character_color=KEN_RED,
                ennemy_color=KEN_GREEN,
                only_punch=os.getenv("TEST_MODE", False),
                temperature=self.temperature,
                sleepy=False,
                model=self.model,
                player_nb=1,
                socket_config=socket_config,
                serving_method=serving_method,
            )
        print(f"[red] Player 1 using: {self.model}")
        self.verify_provider_name()


class Player2(Player):
    def __init__(
        self,
        nickname: str,
        model: str,
        robot_type: str = "text",
        temperature: float = 0.7,
        delay: float = 0.0,
        host: str = "localhost",
        port: int = 38001,
        max_client: int = 1,
        serving_method: str = "remote",
    ):
        self.nickname = nickname
        self.model = model
        self.robot_type = robot_type
        self.temperature = temperature
        self.delay = delay
        if robot_type == "vision":
            self.robot = VisionRobot(
                action_space=None,
                character="Ken",
                side=1,
                model=self.model,
                character_color=KEN_GREEN,
                ennemy_color=KEN_RED,
                temperature=self.temperature,
                sleepy=os.getenv("TEST_MODE", False),
                player_nb=2,
                delay=self.delay,
            )
        elif robot_type == "text":
            self.robot = TextRobot(
                action_space=None,
                character="Ken",
                side=1,
                character_color=KEN_GREEN,
                ennemy_color=KEN_RED,
                temperature=self.temperature,
                sleepy=os.getenv("TEST_MODE", False),
                model=self.model,
                player_nb=2,
                delay=self.delay,
            )
        elif robot_type == "remote":
            socket_config = SocketConfig(
                Host=host,
                Port=port,
                max_client=max_client,
            )
            self.robot = TextLocalRobot(
                action_space=None,
                character="Ken",
                side=1,
                character_color=KEN_GREEN,
                ennemy_color=KEN_RED,
                temperature=self.temperature,
                sleepy=os.getenv("TEST_MODE", False),
                model=self.model,
                player_nb=2,
                delay=self.delay,
                socket_config = socket_config,
                serving_method=serving_method,
            )
        print(f"[green] Player 2 using: {self.model}")
        self.verify_provider_name()


class Episode:
    player_1: Optional[Player1]
    player_2: Player2
    player_1_won: Optional[bool] = None

    def __init__(self, player_1: Optional[Player1], player_2: Player2):
        self.player_1 = player_1
        self.player_2 = player_2

    def save(self):
        if self.player_1 is None:
            # Playing with the controller
            # Results are not saved
            return

        # Write the results to an existing csv with headers "player_1", "player_2", "winner"
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

        # Verifty if the file exists
        if not os.path.exists("results.csv"):
            with open("results.csv", "w") as f:
                f.write(
                    "id,player_1_model,player_1_robot_type, player_1_temperature, player_2_model,player_2_robot_type,player_2_temperature, player_1_won\n"
                )

        with open("results.csv", "a") as f:
            f.write(
                f"{timestamp},{self.player_1.model},{self.player_1.robot_type},{self.player_1.temperature},"
                + f"{self.player_2.model},{self.player_2.robot_type},{self.player_2.temperature},{self.player_1_won}\n"
            )


class Game:
    player_1: Optional[Player1] = None  # First player. None if Human
    player_2: Player2

    render: Optional[bool] = False
    splash_screen: Optional[bool] = False
    save_game: Optional[bool] = False
    characters: Optional[List[str]] = ["Ken", "Ken"]
    outfits: Optional[List[int]] = [1, 3]
    frame_shape: Optional[List[int]] = [0, 0, 0]
    seed: Optional[int] = 42
    settings: EnvironmentSettingsMultiAgent = None  # Settings of the game
    env = None  # Environment of the game

    def __init__(
        self,
        player_1: Optional[Player1],
        player_2: Player2,
        render: bool = False,
        save_game: bool = False,
        splash_screen: bool = False,
        characters: List[str] = ["Ken", "Ken"],
        super_arts: List[int] = [3, 3],
        outfits: List[int] = [1, 3],
        frame_shape: List[int] = [0, 0, 0],
        seed: int = 42,
    ):
        """_summary_

        Args:
            render (bool, optional): Renders the fights. Defaults to False.
            splash_screen (bool, optional): Display the splash screen. Defaults to False.
            characters (List[str], optional): List of the players to have. Defaults to ["Ryu", "Ken"].
            outfits (List[int], optional): Outfits to run. Defaults to [2, 2].
            frame_shape (List[int], optional): Don't know :D . Defaults to [0, 0, 0].
            seed (int, optional): Random seed. Defaults to 42.
        """
        self.render = render
        self.splash_screen = splash_screen
        self.save_game = save_game
        self.characters = characters
        self.super_arts = super_arts
        self.outfits = outfits
        self.frame_shape = frame_shape
        self.seed = seed
        self.settings = self._init_settings()
        self.env = self._init_env(self.settings)
        self.observation, self.info = self.env.reset(seed=self.seed)
        if player_1 is not None:
            self.player_1 = player_1
        else:
            # If player 1 is not provided, we will use the controller
            # The human player will be able to play with the controller
            from diambra.arena.utils.controller import get_diambra_controller

            self.player_1 = None
            self.controller = get_diambra_controller(
                self.env.unwrapped.get_actions_tuples(),
                # force_configure=True,
            )
            self.controller.start()
        self.player_2 = player_2

    def _init_settings(self) -> EnvironmentSettingsMultiAgent:
        """
        Initializes the settings for the game.
        """
        settings = EnvironmentSettingsMultiAgent(
            render_mode="rgb_array",
            splash_screen=self.splash_screen,
        )

        settings.action_space = (SpaceTypes.DISCRETE, SpaceTypes.DISCRETE)
        settings.characters = self.characters
        settings.outfits = self.outfits
        settings.frame_shape = self.frame_shape
        settings.super_art = self.super_arts

        return settings

    def _init_recorder(self) -> RecordingSettings:
        """
        Initializes the recorder for the game.
        """
        if not self.save_game:
            return None
        # Recording settings in root directory
        root_dir = os.path.dirname(os.path.abspath(__file__))
        game_id = "sfiii3n"
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        recording_settings = RecordingSettings()
        recording_settings.dataset_path = os.path.join(
            root_dir, "diambra/episode_recording", game_id, "-", timestamp
        )
        recording_settings.username = "llm-colosseum"

        return recording_settings

    def _init_env(self, settings: EnvironmentSettingsMultiAgent):
        """
        Initializes the environment for the game.
        """
        render_mode = "human" if self.render else "rgb_array"
        recorder_settings = self._init_recorder()
        if self.save_game:
            return make(
                "sfiii3n",
                settings,
                render_mode=render_mode,
                episode_recording_settings=recorder_settings,
            )
        return make("sfiii3n", settings, render_mode=render_mode)

    def _save(self):
        """
        Save the game state.
        """
        pass

    def _determine_winner(self, episode: Episode):
        p1_health = self.observation["P1"]["health"][0]
        p2_health = self.observation["P2"]["health"][0]
        if p1_health > p2_health:
            episode.player_1_won = True
        elif p2_health > p1_health:
            episode.player_1_won = False
        else:
            return "Draw"

    def run(self,delayed=0.0):
        """
        Runs the game with the given settings.
        """
        try:
            self.actions = {
                "agent_0": 0,
                "agent_1": 0,
            }
            self.reward = 0.0

            if self.player_1 is not None:
                self.player_1.robot.observe(self.observation, {}, 0.0)

            self.player_2.robot.observe(self.observation, {}, 0.0)
            # Initialize the episode
            episode = Episode(player_1=self.player_1, player_2=self.player_2)
            # Start the threads that make API calls
            if self.player_1 is not None:
                player1_thread = PlanAndActPlayer1(game=self, episode=episode, serving_type=self.player_1.robot.serving_method)
                player1_thread.start()
            player2_thread = PlanAndActPlayer2(game=self, episode=episode,delayed=delayed, serving_type=self.player_2.robot.serving_method)
            player2_thread.start()
            logger.info(
                f"Game started between {self.player_1.nickname} and {self.player_2.nickname}"
            )
            print("------game simluation starting----------")
            while True:
                # Render the game
                if self.render:
                    self.env.render()

                actions = self.actions
                time.sleep(0.001)
                if self.player_1 is None:
                    # If player 1 is not provided, we use the controller
                    try:
                        # On MacOS we need to install pyobjc
                        # pip install pyobjc
                        # https://stackoverflow.com/questions/76434535/attributeerror-super-object-has-no-attribute-init

                        controller_actions = self.controller.get_actions()
                        actions["agent_1"] = (
                            controller_actions[0] + controller_actions[1]
                        )
                    except Exception as e:
                        print(e)

                if "agent_0" not in actions:
                    actions["agent_0"] = 0
                if "agent_1" not in actions:
                    actions["agent_1"] = 0
                observation, reward, terminated, truncated, info = self.env.step(
                    actions
                )
                # Remove the actions that were executed
                if "agent_0" in self.actions:
                    del actions["agent_0"]
                if "agent_1" in self.actions:
                    del actions["agent_1"]

                self.observation = observation
                self.reward += reward

                p1_wins = observation["P1"]["wins"][0]
                p2_wins = observation["P2"]["wins"][0]

                if p1_wins == 1 or p2_wins == 1:
                    if self.player_1 is not None:
                        player1_thread.running = False
                    player2_thread.running = False
                    episode.player_1_won = p1_wins == 1
                    if episode.player_1_won:
                        print(
                            f"[red] Player1 {self.player_1.robot.model} '{self.player_1.nickname}' won!"
                        )
                        logger.info(f"[red]Player1 {self.player_1.robot.model} '{self.player_1.nickname}' won!")
                    else:
                        print(
                            f"[green] Player2 {self.player_2.robot.model} {self.player_2.nickname} won!"
                        )
                        logger.info(f"[green] Player2 {self.player_2.robot.model} {self.player_2.nickname} won!")
                    episode.save()
                    self.env.close()
                    ##TODO: Replace the line bellow by pass
                    player1_time_list = player1_thread.join()
                    player2_time_list = player2_thread.join()
                    # print(f"Player 1 time list: {player1_time_list}")
                    # print(f"Player 1 average time: {sum(player1_time_list) / len(player1_time_list)}")
                    # print(f"Player 2 time list: {player2_time_list}")
                    # print(f"Player 2 average time: {sum(player2_time_list) / len(player2_time_list)}")
                    return episode.player_1_won
        except Exception as e:
            # self.env.close()
            print(f"Exception: {e}")
            traceback.print_exception(limit=10)
            traceback.print_tb(limit=40)
            if self.player_1 is None:
                self.controller.stop()
            self.env.close()
        try:
            if self.player_1 is None:
                self.controller.stop()
            self.env.close()
        except Exception as e:
            pass  # Ignore the exception
        return 0


class PlanAndAct(Thread):
    def __init__(self, game: Game, episode: Episode, delayed:float = 0.0, serving_type = None):
        self.running = True
        self.game = game
        self.episode = episode
        self.delayed = delayed
        self.serving_type = serving_type

        Thread.__init__(self, daemon=True)
        # atexit.register(self.stop)


class PlanAndActPlayer1(PlanAndAct):
    def run(self) -> None:
        player_1 = []
        connect_flag = False
        while True:
            if self.running:
                if connect_flag is False and self.serving_type == "remote":
                    
                    self.game.player_1.robot.connect_socket()
                    connect_flag = True
 
                if "agent_0" not in self.game.actions:
                    # Plan
                    # print("start acting Player 1")
                    start = time.time()
                    self.game.player_1.robot.plan()
                    # Act
                    self.game.actions["agent_0"] = self.game.player_1.robot.act()
                    # Observe the environment
                    self.game.player_1.robot.observe(
                        self.game.observation, self.game.actions, self.game.reward
                    )
                    end = time.time()
                    period_time = end - start
                    player_1.append(period_time)
            else:
                # return player_1
                print(f"Player 1 time list: {player_1}")
                print(f"Player 1 average time: {sum(player_1) / len(player_1)}")
                logger.info(
                    f"Player 1 time list: {player_1}"
                )
                logger.info(
                    f"Player 1 average time: {sum(player_1) / len(player_1)}"
                )
                break


class PlanAndActPlayer2(PlanAndAct):
    def run(self) -> None:
        player_2 = []
        connect_flag = False
        while True:
            if self.running:
                if connect_flag is False and self.serving_type == "remote":
                    self.game.player_2.robot.connect_socket()
                    connect_flag = True

                if "agent_1" not in self.game.actions:
                    
                    start = time.time()
                    # Plan
                    self.game.player_2.robot.plan()
                    # Act
                    self.game.actions["agent_1"] = self.game.player_2.robot.act()
                    # Observe the environment
                    self.game.player_2.robot.observe(
                        self.game.observation, self.game.actions, -self.game.reward
                    )
                    end = time.time()
                    period_time = end - start
                    player_2.append(period_time)
                    time.sleep(self.delayed * period_time)
            else:
                # return player_2
                print(f"Player 2 time list: {player_2}")
                print(f"Player 2 average time: {sum(player_2) / len(player_2)}")
                logger.info(
                    f"Player 2 time list: {player_2}"
                )
                logger.info(
                    f"Player 2 average time: {sum(player_2) / len(player_2)}"
                )
                break