import os
import asyncio
import websockets
import asyncudp

UDP_IP = os.getenv("UDP_IP", "127.0.0.1")
UDP_PORT = int(os.getenv("UDP_PORT", 12345))

async def websocket_handler(websocket):
    udp_socket = await asyncudp.create_socket(remote_addr=(UDP_IP, UDP_PORT))
    print("Received new connection from", websocket.remote_address)

    async def forward_to_websocket(udp_socket, websocket):
        while True:
            data, addr = await udp_socket.recvfrom()
            await websocket.send(data)

    task1 = asyncio.create_task(forward_to_websocket(udp_socket, websocket))

    try:
        async for message in websocket:
            udp_socket.sendto(message)
    except websockets.exceptions.ConnectionClosed:
        pass
    finally:
        print("Connection to", websocket.remote_address, "is closed")
        task1.cancel()
        udp_socket.close()

async def main():
    async with websockets.serve(websocket_handler, "0.0.0.0", 80):
        print("WebSocket server started")
        while True:
            await asyncio.sleep(3600)

asyncio.run(main())