import asyncio
import json
import signal
import typer
from websockets.asyncio.server import serve, broadcast
from alm_broadcast_server import client

app = typer.Typer()

shutdown_event = asyncio.Event()
connections = set()

ANSI_COLOURS = ["\033[0;31m", "\033[0;32m","\033[0;33m", "\033[0;34m", 
                "\033[0;35m", "\033[0;36m", "\033[0;37m", "\033[0;30m"]


@app.command()
def start(port: int=8001):
    """
    Start a new websocket server, on port 8001 by default.

    The port number to use can be specified with --port.
    Type `exit` or `stop` to shutdown the server.
    """
    asyncio.run(main_server(port))


@app.command()
def connect(port: int=8001):
    """
    Connect to a existing websocket server, on port 8001 by default.

    The port number to use can be specified with --port.
    Type `/exit` to exit the server.
    """
    client.run_client(port)


async def handler(websocket):
    """Handles a new client connection"""
    # username and colour for the new user
    username = f"User{len(connections)+1}"
    colour = ANSI_COLOURS[len(connections) % len(ANSI_COLOURS)]
    
    await websocket.send(f"{username},{colour}")
    connections.add(websocket)

    # new client connected broadcast
    connect_msg = { "message": f"{username} Connected", "username": "server" , "colour": colour}
    broadcast(connections, json.dumps(connect_msg))

    try:
        async for data in websocket:
            broadcast(connections, data)
    finally:
        connections.remove(websocket)
        # client disconnected broadcast
        disconnect_msg = { "message": f"{username} Disconnected", "username": "server", "colour": colour }
        broadcast(connections, json.dumps(disconnect_msg))


async def server_input():
    """Handles user input for shutting down the server"""
    loop = asyncio.get_running_loop()
    try: 
        while True:
            input_text = await loop.run_in_executor(None, input, ">")
            input_text = input_text.strip().lower()
            if (input_text.strip().lower() in ["exit", "stop"]):
                shutdown_event.set()
                return
            elif (input_text == "help"):
                print("\033[1mOptions:\033[0m\n`exit` and `stop` - shuts down the server" \
                      "\n`help` - displayes this message")
    except Exception:
        print("Input error")


def signal_handler():
    """Handles interrupt signals"""
    shutdown_event.set()


async def main_server(port):
    """Starts a websocket server"""
    # Adding a signal handler for (keyboard) interrupts 
    loop = asyncio.get_running_loop()
    loop.add_signal_handler(signal.SIGINT, signal_handler) 

    async with serve(handler, "", port) as server:
        input_task = asyncio.create_task(server_input())
        
        await shutdown_event.wait()

        if not input_task.done():
            input_task.cancel()
            try:
                await input_task
            except asyncio.CancelledError:
                pass
        
        server.close() # closes all client connections then shuts down the server