# online_env.py
import json
import logging
import requests
import websockets
import ssl
import asyncio
import warnings

import urllib3
from urllib.parse import urlencode
from urllib3.exceptions import InsecureRequestWarning
from typing import Optional

import textarena as ta

# 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())
