# online_env.py
import json
import logging
import requests
import websockets
import ssl
import asyncio
import warnings
from urllib.parse import urlencode
from urllib3.exceptions import InsecureRequestWarning

import textarena as ta  # assumes your library defines ta.Env and agents

# Suppress SSL warnings
warnings.filterwarnings("ignore", category=InsecureRequestWarning)

# Constants
WS_SERVER_URI = "wss://api.textarena.ai/ws"
HTTP_SERVER_URI = "https://api.textarena.ai"

NAME_TO_ID_DICT = {
    "Chess-v0": 0,
    "ConnectFour-v0": 1,
    "Debate-v0": 2,
    "DontSayIt-v0": 3,
    "IteratedPrisonersDilemma-v0": 4,
    "Battleship-v0": 5,
    "LiarsDice-v0": 6,
    "Mastermind-v0": 7,
    "Negotiation-v0": 8,
    "Poker-v0": 9,
    "SpellingBee-v0": 10,
    "SpiteAndMalice-v0": 11,
    "Stratego-v0": 12,
    "Tak-v0": 13,
    "TruthAndDeception-v0": 14,
    "UltimateTicTacToe-v0": 15,
    "WordChains-v0": 16,
}


class OnlineEnvWrapper(ta.Env):
    def __init__(self, env_ids: list[int], model_name: str, model_token: str):
        self.env_ids = env_ids
        self.model_name = model_name
        self.model_token = model_token
        self.websocket = None
        self.full_observations: dict[int, list[tuple[int, str]]] = {}
        self.game_over = False
        self.queue_timeout = 1800  # 30 minutes
        self.game_timeout = 300    # 5 minutes
        self.message_queue = asyncio.Queue()
        self.in_game = False

        # State for wrapper compatibility (using a dummy state)
        DummyState = type("DummyState", (), {})  
        self.state = DummyState()
        self.state.role_mapping = {0: "Player 0", 1: "Player 1", -1: "GAME"}

    async def _message_receiver(self):
        """Background task to receive messages."""
        try:
            while True:
                try:
                    message = await self.websocket.recv()
                    await self.message_queue.put(message)
                except websockets.exceptions.ConnectionClosedOK:
                    logging.info("Websocket closed normally")
                    break
                except Exception as e:
                    logging.error(f"Error receiving message: {e}")
                    break
        except Exception as e:
            if not self.game_over:
                logging.error(f"Message receiver error: {e}")

    async def connect(self):
        """Connect to server and queue for game."""
        # Create SSL context
        ssl_context = ssl.create_default_context()
        ssl_context.check_hostname = False
        ssl_context.verify_mode = ssl.CERT_NONE

        # Connect with parameters
        params = {"model_name": self.model_name, "model_token": self.model_token}
        query_params = "?" + urlencode(params)

        self.websocket = await websockets.connect(
            WS_SERVER_URI + query_params,
            ssl=ssl_context,
            ping_interval=20,
            ping_timeout=60,
        )
        logging.info("Connected to server")

        # Start background message receiver
        asyncio.create_task(self._message_receiver())

        # Queue for game
        queue_command = {"command": "queue", "environments": self.env_ids}
        await self.websocket.send(json.dumps(queue_command))
        logging.info("Sent queue request")

        # Wait for queue confirmation
        message = await self.message_queue.get()
        logging.info(f"Queue response: {message}")

    async def get_observation(self) -> tuple[Optional[int], str]:
        """Get next observation."""
        try:
            timeout = self.queue_timeout if not self.in_game else self.game_timeout
            message = await asyncio.wait_for(self.message_queue.get(), timeout=timeout)
            payload = json.loads(message)
            command = payload.get("command")

            if command == "match_found":
                self.in_game = True  # Game started
                player_id = payload.get("player_id")
                obs = payload.get("observation", [])

                if player_id not in self.full_observations:
                    self.full_observations[player_id] = []

                if obs:  # Starting player
                    self.full_observations[player_id].extend(obs)
                    return player_id, self._convert_obs_to_str(player_id)
                else:
                    return None, ""

            elif command == "observation":
                obs = payload.get("observation", [])
                player_id = payload.get("player_id")

                if player_id is not None:
                    if player_id not in self.full_observations:
                        self.full_observations[player_id] = []
                    self.full_observations[player_id].extend(obs)
                    return player_id, self._convert_obs_to_str(player_id)

                return None, ""

            elif command == "game_over":
                self.game_over = True
                return None, ""

            else:
                logging.warning(f"Unknown command: {command}")
                return None, ""

        except asyncio.TimeoutError:
            state = "matchmaking" if not self.in_game else "game"
            logging.error(f"Timeout during {state}")
            self.game_over = True
            return None, ""

        except Exception as e:
            logging.error(f"Error in get_observation: {e}")
            self.game_over = True
            return None, ""

    def _convert_obs_to_str(self, player_id: int) -> str:
        """Format observations for the given player."""
        str_observation = ""
        if player_id in self.full_observations:
            for sender_id, message in self.full_observations[player_id]:
                sender_name = "GAME" if sender_id == -1 else self.state.role_mapping.get(sender_id, f"Player {sender_id}")
                str_observation += f"\n[{sender_name}] {message}"
        recipient_name = self.state.role_mapping.get(player_id, f"Player {player_id}")
        str_observation += f"\n[{recipient_name}]"
        return str_observation

    async def step(self, action: str):
        """Take an action."""
        action_msg = {"command": "action", "action": action}
        await self.websocket.send(json.dumps(action_msg))
        logging.info(f"Sent action: {action}")
        return self.game_over, {}

    async def reset(self):
        """Required by Env; resets the environment."""
        _, observation = await self.get_observation()
        return observation

    async def close(self):
        """Clean up the connection."""
        if self.websocket and not self.websocket.closed:
            await self.websocket.close()
        return {}

def make_online(
    env_ids: list[str],
    model_name: str,
    model_token: Optional[str] = None,
    model_description: Optional[str] = None,
    email: Optional[str] = None,
) -> ta.Env:
    """Create environment wrapper."""
    try:
        env_ids_int = [NAME_TO_ID_DICT[env] for env in env_ids]
    except KeyError as e:
        logging.error(f"Environment {e} not recognized")
        raise

    if model_token is None:
        if model_description is None or email is None:
            raise ValueError("Provide model_description and email if model_token is not given")
        # Register model
        url = f"{HTTP_SERVER_URI}/register_model"
        data = {"model_name": model_name, "description": model_description, "email": email}
        response = requests.post(url, json=data, timeout=10, verify=False)
        response.raise_for_status()
        model_token = response.json()["model_token"]
        logging.info("Model registered successfully")

    return OnlineEnvWrapper(env_ids_int, model_name, model_token)

async def run_game_loop(env, agent):
    """Run a single game loop asynchronously."""
    try:
        await env.connect()
        done = False
        while not done:
            player_id, observation = await env.get_observation()
            if env.game_over:
                break

            if player_id is not None:
                action = agent(observation)
                done, info = await env.step(action)
            else:
                await asyncio.sleep(0.1)

        await env.close()

    except Exception as e:
        logging.error(f"Error in game loop: {e}")
        await env.close()


class SyncOnlineEnv:
    """
    Synchronous wrapper for the asynchronous OnlineEnvWrapper.
    Clients can use reset(), step(), get_observation(), and close() in a blocking fashion.
    """
    def __init__(self, async_env: OnlineEnvWrapper):
        self.async_env = async_env
        self.loop = asyncio.get_event_loop()

    def reset(self):
        return self.loop.run_until_complete(self.async_env.reset())

    def step(self, action: str):
        return self.loop.run_until_complete(self.async_env.step(action))

    def get_observation(self):
        return self.loop.run_until_complete(self.async_env.get_observation())

    def close(self):
        return self.loop.run_until_complete(self.async_env.close())
