from datetime import datetime
from json.decoder import JSONDecodeError
from typing import Dict

from attrs import define
from dragonion_core.proto.web.webmessage import (
    WebBroadcastableMessage,
    WebConnectionMessage,
    WebConnectionResultMessage,
    WebDisconnectMessage,
    WebErrorMessage,
    WebMessageMessage,
    WebNotificationMessage,
    set_time,
    webmessage_error_message_literal,
    webmessages_union,
)
from fastapi import WebSocket

from .connection import Connection
from .exceptions import GotInvalidWebmessage


@define
class Room(object):
    connections: Dict[str, Connection] = {}

    async def accept_connection(self, ws: WebSocket) -> Connection | None:
        """
        Accepts connection, checks username availability and adds it to dict of
        connections
        :param ws: Websocket of connection
        :return:
        """
        print("Incoming connection")
        await ws.accept()
        try:
            connection_message = WebConnectionMessage.from_json(await ws.receive_text())
        except JSONDecodeError:
            await ws.send_text(
                set_time(WebErrorMessage("invalid_webmessage")).to_json()
            )
            await ws.close(reason="invalid_webmessage")
            return

        connection = Connection(
            username=connection_message.username,
            ws=ws,
            public_key=connection_message.public_key,
            password=connection_message.password,
        )

        if connection_message.username in self.connections.keys():
            await connection.send_error("username_exists")
            await ws.close(reason="username_exists")
            return

        self.connections[connection_message.username] = connection
        await connection.send_webmessage(
            WebConnectionResultMessage(
                connected_users=dict(
                    map(
                        lambda i, j: (i, j),
                        [
                            _username
                            for _username in list(self.connections.keys())
                            if self.connections[_username].password
                            == connection_message.password
                        ],
                        [
                            _connection.public_key
                            for _connection in self.connections.values()
                            if _connection.password == connection_message.password
                        ],
                    )
                )
            )
        )

        await self.broadcast_webmessage(connection_message)
        print(f"[{datetime.now().time()}] Accepted {connection_message.username}")
        return connection

    async def broadcast_webmessage(self, obj: webmessages_union):
        """
        Broadcasts WebMessages to all connections in room
        :param obj:
        :return:
        """
        for connection in self.connections.values():
            await connection.send_webmessage(obj)

    async def broadcast_message(self, broadcastable: WebBroadcastableMessage):
        """
        Broadcasts message to every user in room
        :param broadcastable: String object with json representation of
                                   WebBroadcastableMessage
        :return:
        """
        try:
            for to_username in broadcastable.messages.keys():
                try:
                    await self.connections[to_username].send_webmessage(
                        broadcastable.messages[to_username]
                    )
                except KeyError:
                    continue
        except JSONDecodeError:
            raise GotInvalidWebmessage

    async def broadcast_notification(self, message: str):
        """
        Broadcasts notification from server
        :param message: Content
        :return:
        """
        await self.broadcast_webmessage(WebNotificationMessage(message=message))

    async def broadcast_error(self, error_message: webmessage_error_message_literal):
        """
        Broadcasts server error
        :param error_message: See webmessage_error_message_literal
        :return:
        """
        await self.broadcast_webmessage(WebErrorMessage(error_message=error_message))

    async def broadcast_user_disconnected(self, username: str):
        """
        Broadcasts that user is disconnected
        :param username: Username of user that disconnected
        :return:
        """
        await self.broadcast_webmessage(WebDisconnectMessage(username=username))

    async def get_connection_by(self, attribute: str, value: str) -> Connection | None:
        """
        Search for connection by attribute and value in it
        :param attribute:
        :param value:
        :return:
        """
        for connection in self.connections.values():
            if getattr(connection, attribute) == value:
                return connection

    async def disconnect(self, connection: Connection, close_reason: str | None = None):
        """
        Disconnects by connection object.
        :param connection: Object of connection.
                           It can be obtained using get_connection_by
        :param close_reason: Reason if exists
        :return:
        """
        if connection not in self.connections.values():
            return

        del self.connections[connection.username]

        try:
            await connection.ws.close(reason=close_reason)
        except Exception as e:
            assert e

        return connection.username