import asyncio
import websockets
import random
import string
import sqlite3

# Database setup (using SQLite for simplicity)
conn = sqlite3.connect('pairing_database.db')
cursor = conn.cursor()
cursor.execute('CREATE TABLE IF NOT EXISTS pairs (key TEXT, client_ws TEXT)')
conn.commit()

async def handle_client(websocket, path):
    try:
        # Receive the client's request (create a pair or connect to a key)
        request_type = await websocket.recv()

        if request_type == "create_pair":
            # Generate a unique key
            unique_key = generate_key()

            # Save the key and associate it with the current client in the database
            cursor.execute('INSERT INTO pairs VALUES (?, ?)', (unique_key, str(websocket)))
            conn.commit()

            # Inform the client about the created key
            await websocket.send(f"Pair key created: {unique_key}")

        elif request_type == "connect_to_key":
            # Receive the key from the client
            key_to_connect = await websocket.recv()

            # Retrieve all clients associated with the key from the database
            cursor.execute('SELECT client_ws FROM pairs WHERE key=?', (key_to_connect,))
            results = cursor.fetchall()

            if results:
                # Pair the current client with all clients associated with the key
                for result in results:
                    other_client_ws = result[0]
                    if other_client_ws != str(websocket):
                        await websocket.send(f"Connected to key {key_to_connect}. You are paired with {other_client_ws}")
                        await other_client_ws.send(f"Connected to key {key_to_connect}. You are paired with {str(websocket)}")
            else:
                # Inform the client that the key does not exist
                await websocket.send(f"Key {key_to_connect} does not exist.")

    except websockets.exceptions.ConnectionClosed:
        pass  # Connection closed, handle appropriately

def generate_key():
    # Generate a random key (you might want to implement a more robust key generation logic)
    return ''.join(random.choices(string.ascii_uppercase + string.digits, k=8))

start_server = websockets.serve(handle_client, "localhost", 8765)

async def main():
    print("WebSocket server is running. Waiting for connections...")
    await start_server
    await asyncio.Event().wait()

if __name__ == "__main__":
    asyncio.run(main())