import struct
from rich.console import Console
import asyncio


class Server:
    BUFFER_SIZE = 1024
    NOTIFICATION_DIR = 'notifications'
    console = Console()

    def __init__(self, ip: str, port: int, producer_port: int, notification_dir: str = 'notifications'):
        self.NOTIFICATION_DIR = notification_dir
        self.__client_socket = None
        self.__producer_socket = None
        self.__ip = ip
        self.__port = port
        self.__producer_port = producer_port
        self.__client_writers = []

    async def up(self):
        self.__client_socket = await asyncio.start_server(
            self.handle_client, self.__ip, self.__port
        )

        self.__producer_socket = await asyncio.start_server(
            lambda reader, writer: self.handle_producer(reader, writer, self.NOTIFICATION_DIR),
            self.__ip, self.__producer_port
        )

        self.console.print(f'Serves on {self.__ip}:{self.__port} for clients.')
        self.console.print(f'Receives notifications from producers at {self.__ip}:{self.__producer_port}.')

        async with self.__producer_socket, self.__client_socket:
            await asyncio.gather(
                self.__producer_socket.serve_forever(),
                self.__client_socket.serve_forever()
            )

    @staticmethod
    async def send_message(writer, message: str):
        writer.write(struct.pack('<L', len(message)))
        writer.write(message.encode('utf-8'))

    @staticmethod
    async def receive_message(reader) -> str:
        size, = struct.unpack('<L', await reader.readexactly(4))
        message = await reader.readexactly(size)
        return message.decode('utf-8')

    async def broadcast_message(self, message):
        for writer in self.__client_writers:
            await self.send_message(writer, message)

    async def handle_producer(self, reader, writer, notification_dir):
        addr = writer.get_extra_info('peername')
        self.console.print(f"A producer ({addr}) connected.")

        try:
            while True:
                message = await self.receive_message(reader)
                print(f"{message=}")
                await self.broadcast_message(message)

        except asyncio.CancelledError:
            pass
        finally:
            self.console.print(f"The producer ({addr}) disconnected.")
            writer.close()

    async def handle_client(self, reader, writer):
        addr = writer.get_extra_info('peername')
        self.console.print(f"A client ({addr}) connected.")
        self.__client_writers.append(writer)

        try:
            while True:
                pass
        except asyncio.CancelledError:
            pass
        finally:
            self.console.print(f"The client ({addr}) disconnected.")
            self.__client_writers.remove(writer)
            writer.close()