from websockets import *
import asyncio, json, prompt_toolkit

class TestClient:
    def __init__(self, address: str = "127.0.0.1", port: int = 8765) -> None:
        self.connection: ClientConnection = None
        self.msgId: int = 0
        self.connected = True
        asyncio.run(self.main(address, port))
    
    async def main(self, address: str, port: int) -> None:
        await self.connect(address, port)
        tasks: list[asyncio.Task] = []
        async with asyncio.TaskGroup() as tg:
            tasks.append(tg.create_task(self.__message__()))
            tasks.append(tg.create_task(self.receive()))
        print("Finished")
        await self.connection.close()

    async def connect(self, address: str = "127.0.0.1", port: int = 8765) -> None:
        uri: str = f"ws://{address}:{port}"
        self.connection = await connect(uri)

    async def __message__(self) -> None:
        session = prompt_toolkit.PromptSession()
        while (self.connected):
            text: str = await session.prompt_async()
            print(f"Text: '{text}'")
            diction: dict = {"ID": self.msgId, "message": text}
            msg = json.dumps(diction)
            await self.connection.send(msg)
            self.msgId += 1
            print(f"Client sent {msg}")
    
    async def receive(self) -> None:
        while (self.connected):
            response: str = await self.connection.recv()
            print(f"Received: {response}")

if __name__ == "__main__":
    client = TestClient()