import asyncio
import websockets
import sqlite3
import secrets

# SQLite database setup
conn = sqlite3.connect('pairing.db')
cursor = conn.cursor()

# Create Pairing table if not exists
cursor.execute('''
    CREATE TABLE IF NOT EXISTS Pairing (
        key TEXT PRIMARY KEY,
        client_address TEXT
    )
''')
conn.commit()

async def handle_client(websocket, path):
    try:
        # Receive a command from the client
        command = await websocket.recv()

        if command == "create_pair":
            # Create a unique key
            key = secrets.token_hex(16)

            # Save the key and client address in the database
            cursor.execute('INSERT INTO Pairing (key, client_address) VALUES (?, ?)', (key, str(websocket.remote_address)))
            conn.commit()

            # Inform the client about the generated key
            await websocket.send(f"Your pair key is: {key}")

        elif command.startswith("connect_to_key"):
            # Extract the key from the command
            key_to_connect = command.split(" ")[1]

            # Check if the key exists in the database
            cursor.execute('SELECT client_address FROM Pairing WHERE key = ?', (key_to_connect,))
            result = cursor.fetchone()

            if result:
                other_client_address = result[0]

                # Inform both clients about the successful pairing
                await websocket.send(f"Connected to {other_client_address}")
                await websocket.send(f"You are paired with {other_client_address}")

                # Create a bidirectional bridge
                other_client = next(c for c in connected_clients if str(c.remote_address) == other_client_address)
                asyncio.ensure_future(bridge_clients(websocket, other_client))
                asyncio.ensure_future(bridge_clients(other_client, websocket))

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

async def bridge_clients(source, destination):
    try:
        while True:
            # Receive data from the source client
            data_from_source = await source.recv()

            # Forward the data to the destination client
            await destination.send(data_from_source)

    except websockets.exceptions.ConnectionClosed:
        pass  # Exit the loop if connection is closed

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__":
    connected_clients = set()
    asyncio.run(main())