#!python

# WS server example that synchronizes state across clients

import asyncio
import json
import logging
import websockets

import oneiot_core.env as env

logging.basicConfig()

USERS = set()

async def notify_message(message):
    if USERS:  # asyncio.wait doesn't accept an empty list
        await asyncio.wait([user.send(message) for user in USERS])


async def notify_connect(websocket):
    if USERS:  # asyncio.wait doesn't accept an empty list
        message = {
            'id': 'core.device.connect',
            'data': {
                'address': {
                    'ip': websocket.remote_address[0],
                    'socket': websocket.remote_address[1]
                }
            }
        }
        message = json.dumps(message)
        await asyncio.wait([user.send(message) for user in USERS])


async def notify_disconnect(websocket):
    if USERS:  # asyncio.wait doesn't accept an empty list
        message = {
            'id': 'core.device.disconnect',
            'data': {
                'address': {
                    'ip': websocket.remote_address[0],
                    'socket': websocket.remote_address[1]
                }
            }
        }
        message = json.dumps(message)
        await asyncio.wait([user.send(message) for user in USERS])

async def notify_error(websocket, transmission):
    if USERS:  # asyncio.wait doesn't accept an empty list
        message = {
            'id': 'core.error.transmissionError',
            'data': {
                'address': {
                    'ip': websocket.remote_address[0],
                    'socket': websocket.remote_address[1]
                },
                'transmission': json.dumps(transmission)
            }
        }
        message = json.dumps(message)
        await asyncio.wait([user.send(message) for user in USERS])

async def register(websocket):
    USERS.add(websocket)
    await notify_connect(websocket)


async def unregister(websocket):
    USERS.remove(websocket)
    await notify_disconnect(websocket)


async def counter(websocket, path):
    # register(websocket) sends user_event() to websocket
    await register(websocket)
    try:
        async for message in websocket:
            try:
                data = json.loads(message)
            except:
                data = message
            if isinstance(data, dict) and "id" in data:
                cleaned_data = {'id': data['id']}
                if "data" in data:
                    cleaned_data['data'] = data['data']
                await notify_message(json.dumps(cleaned_data))
            else:
                await notify_error(websocket, data)
                logging.error(f"unsupported event: {data}")
    finally:
        await unregister(websocket)


start_server = websockets.serve(counter, "0.0.0.0", env.var("ONEIOT_C_PORT"))

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()