import asyncio
import json
import logging
import aiomysql
from websockets.server import serve
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK

# --- CONFIGURATION ---
HOST = "0.0.0.0"
PORT = 8765

# MySQL Config (Matches config.php)
DB_HOST = "127.0.0.1"
DB_NAME = "studio_webrtc"
DB_USER = "airlink"
DB_PASS = "F@lsgrave962"

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

# Global Registry: node_id (int) -> WebSocket connection object
connected_nodes = {}

async def get_db_pool():
    """Create an asynchronous MySQL connection pool."""
    return await aiomysql.create_pool(
        host=DB_HOST,
        port=3306,
        user=DB_USER,
        password=DB_PASS,
        db=DB_NAME,
        autocommit=True
    )

async def validate_session(pool, token):
    """Verify session token against PHP/MySQL database and retrieve Node ID."""
    async with pool.acquire() as conn:
        async with conn.cursor(aiomysql.DictCursor) as cursor:
            query = """
                SELECT id, user_id, node_friendly_name 
                FROM nodes 
                WHERE session_token = %s
            """
            await cursor.execute(query, (token,))
            result = await cursor.fetchone()
            return result

async def update_node_status(pool, node_id, status):
    """Mark node as 'online' or 'offline' in the database."""
    async with pool.acquire() as conn:
        async with conn.cursor() as cursor:
            query = "UPDATE nodes SET status = %s, last_active = NOW() WHERE id = %s"
            await cursor.execute(query, (status, node_id))

async def broadcast_node_list():
    """Notify all connected WebSocket clients when the active node list updates."""
    active_ids = list(connected_nodes.keys())
    payload = json.dumps({
        "type": "presence_update",
        "active_node_ids": active_ids
    })
    
    # Broadcast to all connected WebSockets concurrently
    if connected_nodes:
        await asyncio.gather(
            *[ws.send(payload) for ws in connected_nodes.values()],
            return_exceptions=True
        )

async def handler(websocket, path, pool):
    """Handle incoming WebSocket connections and routing WebRTC signals."""
    authenticated_node_id = None

    try:
        # Step 1: Wait for initial authentication payload from client
        raw_msg = await websocket.recv()
        data = json.loads(raw_msg)

        if data.get("type") != "register":
            await websocket.send(json.dumps({"type": "error", "message": "First message must be registration."}))
            await websocket.close(1008, "Unauthenticated")
            return

        token = data.get("token")
        node_data = await validate_session(pool, token)

        if not node_data:
            await websocket.send(json.dumps({"type": "error", "message": "Invalid or expired session token."}))
            await websocket.close(1008, "Invalid Token")
            return

        authenticated_node_id = node_data["id"]
        
        # Register in-memory session
        connected_nodes[authenticated_node_id] = websocket
        await update_node_status(pool, authenticated_node_id, "online")

        logging.info(f"Node {authenticated_node_id} ('{node_data['node_friendly_name']}') connected successfully.")

        # Confirm registration back to client
        await websocket.send(json.dumps({
            "type": "registered",
            "node_id": authenticated_node_id,
            "node_name": node_data["node_friendly_name"]
        }))

        # Notify network of new online node
        await broadcast_node_list()

        # Step 2: Continuous Signaling Loop
        async for message in websocket:
            try:
                signal = json.loads(message)
                action_type = signal.get("type")
                target_node_id = signal.get("target_node_id")

                # WebRTC Handshake messages: 'offer', 'answer', 'ice-candidate'
                if action_type in ["offer", "answer", "ice-candidate"]:
                    if target_node_id and target_node_id in connected_nodes:
                        target_ws = connected_nodes[target_node_id]
                        
                        # Attach sender details so recipient knows who sent the SDP/ICE signal
                        signal["sender_node_id"] = authenticated_node_id
                        await target_ws.send(json.dumps(signal))
                        logging.debug(f"Relayed {action_type} from Node {authenticated_node_id} -> Node {target_node_id}")
                    else:
                        await websocket.send(json.dumps({
                            "type": "error",
                            "message": f"Target node {target_node_id} is offline or unreachable."
                        }))

                elif action_type == "ping":
                    await update_node_status(pool, authenticated_node_id, "online")
                    await websocket.send(json.dumps({"type": "pong"}))

            except json.JSONDecodeError:
                logging.warning(f"Received malformed JSON from Node {authenticated_node_id}")

    except (ConnectionClosedOK, ConnectionClosedError):
        logging.info(f"Client disconnected: Node {authenticated_node_id}")
    except Exception as e:
        logging.error(f"Unexpected error handling Node {authenticated_node_id}: {e}")
    finally:
        # Step 3: Cleanup on disconnection
        if authenticated_node_id in connected_nodes:
            del connected_nodes[authenticated_node_id]
            await update_node_status(pool, authenticated_node_id, "offline")
            logging.info(f"Cleaned up Node {authenticated_node_id}")
            await broadcast_node_list()

async def main():
    pool = await get_db_pool()
    logging.info(f"Database pool established for {DB_NAME}")

    # Wrap handler to inject database connection pool
    async def bound_handler(websocket, path):
        await handler(websocket, path, pool)

    async with serve(bound_handler, HOST, PORT):
        logging.info(f"Python Signaling Server running on ws://{HOST}:{PORT}")
        await asyncio.Future()  # Run forever

if __name__ == "__main__":
    try:
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main())
    except KeyboardInterrupt:
        logging.info("Signaling Server shutting down.")
    finally:
        loop.close()