import websockets
import asyncio
import uuid
from enum import Enum
import time
from pydantic import BaseModel
import logging
import json
import traceback

from multiplayer_crosswords.crossword import Crossword
from multiplayer_crosswords.crossword_algorithm import Orientation
from multiplayer_crosswords.dictionary import Dictionary, Word
from multiplayer_crosswords.utils import load_en_dictionary, load_de_dictionary

from .server_config import ServerConfig
from .server_utils import Languages
from .websocket_connection_handler import WebsocketConnectionHandler

from . import client_messages
from . import server_messages

logger = logging.getLogger(__name__)

# set global logging level to DEBUG for demonstration purposes
logging.basicConfig(level=logging.DEBUG)

class MultiplayerSession:

    @property
    def crossword(self) -> Crossword:
        self._assure_is_locked()
        return self._crossword

    @property
    def session_id(self) -> str:
        self._assure_is_locked()
        return self._session_id
    
    @property
    def clients(self) -> set:
        self._assure_is_locked()
        return self._clients

    @property
    def lock(self) -> asyncio.Lock:
        self.update_activity()
        return self._lock

    @property
    def is_idle(self) -> bool:
        idle_time = time.monotonic() - self._last_active_time
        return idle_time > self._max_idle_time_seconds


    def _assure_is_locked(self):
        if not self._lock.locked():
            raise RuntimeError("MultiplayerSession methods that modify state must be called while holding the session lock.")

    def __init__(self, session_id: str, crossword: Crossword):
        self._session_id = session_id
        self._crossword = crossword
        self._last_active_time = time.monotonic()
        self._clients = set()
        self._max_idle_time_seconds = ServerConfig.get_config().MAX_SESSION_IDLE_TIME_SECONDS
        self._lock = asyncio.Lock()

    def update_activity(self):
        self._last_active_time = time.monotonic()

    def add_client(self, client: WebsocketConnectionHandler):
        self._assure_is_locked()
        self._clients.add(client)

    def remove_client(self, client: WebsocketConnectionHandler):
        self._assure_is_locked()
        self._clients.discard(client)
    
    def remove_unconnected_clients(self):
        self._assure_is_locked()
        self._clients = {c for c in self._clients if not c.is_closed()}

    async def send_message_to_all_clients(self, message: dict):
        logger.debug("Broadcasting message to all clients in session %s: %s. #clients: %s", self._session_id, message, len(self._clients))
        for client in self._clients.copy():
            if client.is_closed():
                logger.debug("Removing closed client %s from session %s", client, self._session_id)
                self._clients.discard(client)
                continue
            try:
                await client.send(message=message)
                logger.debug("Message sent to client %s: %s", client, message)
            except Exception as e:
                logger.warning("error sending message to client %s: %s", client, str(e))
                logger.warning("stack trace: %s", traceback.format_exc())


class MultiplayerSessionManager(object):

    @classmethod
    def instance(cls) -> "MultiplayerSessionManager":
        if not hasattr(cls, "_singleton_instance"):
            cls._singleton_instance = cls()
        return cls._singleton_instance

    def __init__(self):
        self.__class__._singleton_instance = self
        self._sessions: dict[str, MultiplayerSession] = {}
        self._sessions_lock = asyncio.Lock()
        self._grid_block_ratio = ServerConfig.get_config().GRID_BLOCK_RATIO

    async def remove_idle_sessions(self):
        """Remove all idle sessions that have exceeded their max idle time."""
        async with self._sessions_lock:
            idle_session_ids = [
                session_id for session_id, session in self._sessions.items()
                if session.is_idle
            ]
            for session_id in idle_session_ids:
                logger.info("Removing idle session %s", session_id)
                del self._sessions[session_id]
            if idle_session_ids:
                logger.info("Removed %d idle sessions", len(idle_session_ids))

    async def create_session(self, lang: str | Languages, grid_w: int, grid_h: int) -> MultiplayerSession:
        # Remove idle sessions before creating a new one
        await self.remove_idle_sessions()
        async with self._sessions_lock:
            
            
            if isinstance(lang, str):
                lang = Languages(lang)
            dictionary = lang.load_dictionary()
            max_tries = 4
            for i in range(max_tries):
                crossword = Crossword.generate(
                    dictionary=dictionary,
                    seed=None,
                    grid_width=grid_w,
                    grid_height=grid_h,
                    grid_block_ratio=self._grid_block_ratio,
                )
                if crossword is not None:
                    break
            if crossword is None:
                raise RuntimeError("Failed to generate crossword for the given parameters.")
            session_id = str(uuid.uuid4())
            session = MultiplayerSession(session_id=session_id, crossword=crossword)
            self._sessions[session_id] = session
            return session

class WebsocketCrosswordServer(object):

    @staticmethod
    def ensure_valid_letter(letter: str):
        if len(letter) != 1 and letter != "":
            raise ValueError("Letter must be a single character A-Z or empty string.")
        if letter != "" and (not letter.isalpha()):
            raise ValueError("Letter must be a single character A-Z or empty string.")
        return letter.upper()
    
    @WebsocketConnectionHandler.register_message_handler(
        model=client_messages.RequestAvailableSessionPropertiesClientMessage
    )
    @staticmethod
    async def handle_request_available_session_properties(handler: WebsocketConnectionHandler, message: client_messages.RequestAvailableSessionPropertiesClientMessage):
        server_config = ServerConfig.get_config()
        response = server_messages.AvailableSessionPropertiesServerMessage(
            supported_languages=list(reversed([lang.value for lang in Languages])),
            min_grid_size=server_config.MIN_GRID_SIZE,
            max_grid_size=server_config.MAX_GRID_SIZE,
            board_size_presets={
                preset.name.lower(): preset.to_dimensions()
                for preset in client_messages.BoardSizePreset
            }
        )
        await handler.send(message=response.model_dump())

    @WebsocketConnectionHandler.register_message_handler(
        model=client_messages.NewMultiplayerSessionClientMessage
    )
    @staticmethod
    async def handle_new_multiplayer_session(handler: WebsocketConnectionHandler, message: client_messages.NewMultiplayerSessionClientMessage):
        message.validate()

        try:
            session = await MultiplayerSessionManager.instance().create_session(
                lang=message.lang,
                grid_w=message.grid_w,
                grid_h=message.grid_h,
            )
        except Exception as e:
            error_response = server_messages.ServerErrorMessage(
                error_message=f"Error creating session: {str(e)}"
            )
            await handler.send(message=error_response.model_dump())
            return


        async with session.lock:
            session.add_client(handler)
            response = server_messages.SessionCreatedServerMessage(
                session_id=session.session_id
            )
            await handler.send(message=response.model_dump())
    
    @WebsocketConnectionHandler.register_message_handler(
        model=client_messages.SubscribeSessionClientMessage
    )
    @staticmethod
    async def handle_request_full_session_state(handler: WebsocketConnectionHandler, message: client_messages.SubscribeSessionClientMessage):
        session_id = message.session_id
        session_manager = MultiplayerSessionManager.instance()
        async with session_manager._sessions_lock:
            if session_id not in session_manager._sessions:
                error_response = server_messages.ServerErrorMessage(
                    error_message=f"Session with id {session_id} not found."
                )
                await handler.send(message=error_response.model_dump())
                return
            session = session_manager._sessions[session_id]
        
        async with session.lock:
            grid_state = session.crossword.current_grid
            # make all letters uppercase for the client
            for row in range(len(grid_state)):
                for col in range(len(grid_state[0])):
                    grid_state[row][col] = grid_state[row][col].upper()
            clues_across, clues_down = session.crossword.get_clues()
            clue_positions_across, clue_positions_down = session.crossword.get_clue_positions()

            # get all sovled positions by looping over all words
            solved_positions = set()
            for cw in session.crossword.words:
                if cw.solved:
                    for i in range(len(cw.word)):
                        if cw.orientation == Orientation.HORIZONTAL:
                            row = cw.start_y
                            col = cw.start_x + i
                        else:
                            row = cw.start_y + i
                            col = cw.start_x
                        solved_positions.add((col, row))

            solved_positions = list(solved_positions)

            solution_word_positions = []
            positions = session.crossword.solution_word_positions
            for pos in positions:
                # Convert from (row, col) to (col, row) for client
                row, col = pos
                solution_word_positions.append((col, row))

            response = server_messages.SendFullSessionStateServerMessage(
                session_id=session.session_id,
                grid=grid_state,
                clues_across=clues_across,
                clues_down=clues_down,
                clue_positions_across=clue_positions_across,
                clue_positions_down=clue_positions_down,
                solved_positions=solved_positions,
                solution_word_positions=solution_word_positions,
            )
            # register the client to the session
            session.add_client(handler)
            await handler.send(message=response.model_dump())
    
    @WebsocketConnectionHandler.register_message_handler(
        model=client_messages.UpdateLetterClientMessage
    )
    @staticmethod
    async def handle_update_letter(handler: WebsocketConnectionHandler, message: client_messages.UpdateLetterClientMessage):
        session_id = message.session_id
        session_manager = MultiplayerSessionManager.instance()
        async with session_manager._sessions_lock:
            if session_id not in session_manager._sessions:
                error_response = server_messages.ServerErrorMessage(
                    error_message=f"Session with id {session_id} not found."
                )
                await handler.send(message=error_response.model_dump())
                return
            session = session_manager._sessions[session_id]
        async with session.lock:
            crossword = session.crossword
            msg_letter = WebsocketCrosswordServer.ensure_valid_letter(message.letter)
            current_grid = crossword.current_grid
            if not (0 <= message.row < len(current_grid)) or not (0 <= message.col < len(current_grid[0])):
                error_response = server_messages.ServerErrorMessage(
                    error_message=f"Row or column out of bounds."
                )
                await handler.send(message=error_response.model_dump())
                return
            current_grid_letter = current_grid[message.row][message.col]
            if current_grid_letter.upper() == msg_letter.upper():
                # No change
                return
            # check if the letter already is solved, if so, ignore the update
            words_at_position = crossword.get_words_by_y_x_position(y=message.row, x=message.col)
            if any(cw.solved for cw in words_at_position):
                logger.info("Ignoring update to already solved position (%d, %d) in session %s", message.col, message.row, session.session_id)

                 # send letter again to client to ensure they have the correct letter
                msg = server_messages.LetterUpdateBroadcastServerMessage(
                    session_id=session.session_id,
                    row=message.row,
                    col=message.col,
                    letter=msg_letter.upper(),
                    is_solved=is_solved
                )
                messages = [msg]

            else:
            # also check if the position is 
                crossword.place_letter(
                    x=message.col,
                    y=message.row,
                    letter=msg_letter.lower(),
                )
            
                words_at_position = crossword.get_words_by_y_x_position(y=message.row, x=message.col)
                is_solved = any(cw.solved for cw in words_at_position)
                if is_solved:
                    logger.info("Word solved at position (%d, %d) in session %s", message.col, message.row, session.session_id)
                    messages = []
                    for cw in words_at_position:
                        if cw.solved:
                            logger.info("Solved word: %s", cw.word)
                            # go through each letter in the word and create a message
                            for i in range(len(cw.word)):
                                if cw.orientation == Orientation.HORIZONTAL:
                                    row = cw.start_y
                                    col = cw.start_x + i
                                else:
                                    row = cw.start_y + i
                                    col = cw.start_x
                                letter = cw.word[i].upper()
                                msg = server_messages.LetterUpdateBroadcastServerMessage(
                                    session_id=session.session_id,
                                    row=row,
                                    col=col,
                                    letter=letter,
                                    is_solved=True
                                )
                                messages.append(msg)

                else:
                    msg = server_messages.LetterUpdateBroadcastServerMessage(
                        session_id=session.session_id,
                        row=message.row,
                        col=message.col,
                        letter=msg_letter.upper(),
                        is_solved=is_solved
                    )
                    messages = [msg]
        
        # NOTE: we do this purposefully outside of the session lock to avoid
        # potential deadlocks if sending messages takes time.
        # this could cause clients to receive messages slightly out of order
        for broadcast_message in messages:
            await session.send_message_to_all_clients(message=broadcast_message.model_dump())


    def __init__(self, host: str, port: int):
        self._host = host
        self._port = port
        
    
    def run(self):
        # Newer versions of the `websockets` library call the handler with a
        # single `connection` argument (the WebSocketServerProtocol). Accept
        # that and pass it into our connection handler wrapper.
        async def connection_handler(connection: websockets.WebSocketServerProtocol):
            handler = WebsocketConnectionHandler(websocket=connection)
            await handler.run()

        async def _main():
            logger.info("Starting WebSocket server on %s:%d", self._host, self._port)
            # Use the async context manager for the server; this will keep
            # the server running until cancelled. Using asyncio.run() to
            # start the loop ensures we don't depend on an existing running
            # event loop (avoids "no running event loop" errors).
            async with websockets.serve(
                handler=connection_handler,
                host=self._host,
                port=self._port,
            ):
                # Wait forever (until process is killed or cancelled)
                await asyncio.Event().wait()

        # Start the server with a fresh event loop
        asyncio.run(_main())