import asyncio
import json 
from websockets.asyncio.client import connect
from websockets.exceptions import InvalidHandshake
from prompt_toolkit import PromptSession, print_formatted_text, ANSI
from prompt_toolkit.patch_stdout import patch_stdout

ANSI_RESET = "\033[0m"


async def connect_server(port):
    """Handles the server connection and runs the communication tasks"""
    uri = f"ws://localhost:{port}"

    try: 
        async with connect(uri) as websocket:
            client_data = await websocket.recv()
            client_data = client_data.split(",")
            client_username = client_data[0]
            client_colour = client_data[1]    

            session = PromptSession()

            # running as concurrent tasks 
            send_task = asyncio.create_task(send(websocket, session, client_colour, client_username))
            receive_task = asyncio.create_task(receive(websocket, client_username))
            # running until the first task is completed
            done, pending = await asyncio.wait([receive_task, send_task], return_when=asyncio.FIRST_COMPLETED)
            
            # canceling other task
            for tasks in pending:
                tasks.cancel()
                try:
                    await tasks
                except asyncio.CancelledError:
                    pass
        
        print_formatted_text(ANSI(f"\r\033[1mServer has shutdown{ANSI_RESET}"))
    except (OSError, InvalidHandshake):
        print(f"\033[1mCould not connect to server{ANSI_RESET}")
        pass


async def send(websocket, session, client_colour, client_username):
    """Accepts user input and sends it to the server"""
    try:    
        while True:
            with patch_stdout():
                send_message = await session.prompt_async(">")
            send_message = send_message.strip().lower()

            if (send_message == "/exit"):
                await websocket.close()
                return
            elif (send_message == "/name"):
                print_formatted_text(ANSI(f"\r{client_colour}{client_username}{ANSI_RESET}"))
            elif (send_message == "/help"):
                print_formatted_text(ANSI(f"\r\033[1mOptions:{ANSI_RESET}\n`/name` - returns your username" \
                                           "\n`/exit` - exits the server" \
                                           "\n`/help` - displays this message"))
            elif (send_message != ""):
                data = { "message": send_message, "username": client_username, "colour": client_colour}
                await websocket.send(json.dumps(data))
    except Exception:
        print_formatted_text("\rError sending message")


async def receive(websocket, client_username):
    """Receives and displays incoming messages"""
    try:
        async for data in websocket:
            loaded_data = json.loads(data)

            message = loaded_data["message"]
            name = loaded_data["username"]
            colour = loaded_data["colour"]

            # Broadcasted messages
            if (name == "server"):
                print_formatted_text(ANSI(f"\r{colour}{message}{ANSI_RESET}"))     
            # Client messages
            elif (name != client_username):
                print_formatted_text(ANSI(f"\r{colour}{name}{ANSI_RESET}>{message}"))
    except Exception:
        print_formatted_text("\rError receiving message")


def run_client(port):
    """Entry point"""
    try:
        asyncio.run(connect_server(port))
    except:
        pass