import logging
import json
import traceback

from pydantic import BaseModel
import websockets

from .server_config import ServerConfig
from .server_utils import Languages
from .client_messages import ClientMessageBase
from .server_messages import ServerMessageBase, ServerErrorMessage

logger = logging.getLogger(__name__)


class WebsocketConnectionHandler:

    message_handlers: dict[str, callable] = {}
    message_models: dict[str, type[BaseModel]] = {}

    @classmethod
    def register_message_handler(cls, model: type[ClientMessageBase]):
        def decorator(func: callable):
            message_type = model.get_type()
            cls.message_handlers[message_type] = func
            cls.message_models[message_type] = model
            return func
        return decorator

    @property
    def websocket(self) -> websockets.WebSocketServerProtocol:
        return self._websocket

    def __init__(self, websocket: websockets.WebSocketServerProtocol):
        self._websocket = websocket
        self._is_closed = False

    def is_closed(self) -> bool:
        """Check if the connection is closed"""
        return self._is_closed

    async def handle_message(self, message: dict):
        try:
            message_type = message.get("type")
            if message_type not in self.message_handlers:
                raise ValueError(f"Unknown message type: {message_type}")
            model_cls = self.message_models[message_type]
            model_instance = model_cls(**message)
            handler = self.message_handlers[message_type]
            await handler(self, model_instance)
        except Exception as e:
            logger.warning("error handling message: %s", str(e))
            logger.warning("stack trace: %s", traceback.format_exc())
            error_response = ServerErrorMessage(
                error_message=f"Error handling message: {str(e)}\n {traceback.format_exc()}"
            )
            await self.send(message=error_response.model_dump())

    async def send(self, message: dict):
        string_message = json.dumps(message)
        logger.debug("sending message: %s", string_message)
        try:
            await self._websocket.send(string_message)
        except Exception as e:
            logger.warning("error sending message: %s", str(e))
            logger.warning("stack trace: %s", traceback.format_exc())
    
    async def close(self):
        try:
            if not self._is_closed:
                self._is_closed = True
                await self._websocket.close()
                logger.debug("WebSocket closed")
        except Exception as e:
            logger.warning("error closing websocket: %s", str(e))
            logger.warning("stack trace: %s", traceback.format_exc())
    
    async def run(self):
        try:
            async for message in self._websocket:
                try:
                    json_message = json.loads(message)
                    await self.handle_message(json_message) 
                except ValueError as e:
                    logger.warning("received unprocessable message %s", str(e))
                    logger.warning("stack trace: %s", traceback.format_exc())
        except Exception as e:
            logger.warning("error in websocket connection: %s", str(e))
            logger.warning("stack trace: %s", traceback.format_exc())
            self._is_closed = True
        finally:
            self._is_closed = True
