#!/usr/bin/env python3
"""Simple WebSocket echo server for testing DPDK WebSocket client."""

import asyncio
import websockets
import json
import sys
import time

connected_clients = set()

async def handler(websocket):
    client_addr = websocket.remote_address
    path = websocket.request.path if hasattr(websocket, 'request') else "/"
    print(f"[CONNECT] {client_addr}, path: {path}")
    connected_clients.add(websocket)
    try:
        async for message in websocket:
            print(f"[RECV] {message}")
            try:
                data = json.loads(message)
                action = data.get("action", "")
                topic = data.get("topic", "")
                if action == "subscribe":
                    resp = json.dumps({"status":"ok","action":"subscribed","topic":topic,"timestamp":int(time.time())})
                    await websocket.send(resp)
                    print(f"[SEND] {resp}")
                    push = json.dumps({"type":"push","topic":topic,"data":f"Welcome to {topic}!","timestamp":int(time.time())})
                    await websocket.send(push)
                    print(f"[SEND] {push}")
                    continue
            except (json.JSONDecodeError, TypeError):
                pass
            response = f"ECHO: {message}"
            await websocket.send(response)
            print(f"[SEND] {response}")
    except websockets.exceptions.ConnectionClosed:
        print(f"[DISCONNECT] {client_addr}")
    finally:
        connected_clients.discard(websocket)

async def push_loop(interval=10):
    count = 0
    while True:
        await asyncio.sleep(interval)
        if connected_clients:
            count += 1
            msg = json.dumps({"type":"push","data":f"Periodic update #{count}","timestamp":int(time.time())})
            websockets.broadcast(connected_clients, msg)
            print(f"[PUSH] {msg}")

async def main(port=8080):
    print(f"WebSocket Test Server on ws://0.0.0.0:{port}")
    async with websockets.serve(handler, "0.0.0.0", port):
        asyncio.create_task(push_loop(10))
        print(f"[READY] Waiting for DPDK client...")
        await asyncio.Future()

if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8080
    asyncio.run(main(port))