import asyncio
import json
import uuid
import websockets

clients = {}
client_to_browsers = {}

WS_MAX_SIZE = 100 * 1024 * 1024


async def send_json(ws, payload):
	try:
		await ws.send(json.dumps(payload))
	except:
		pass


def new_client_id():
	return str(uuid.uuid4())


def add_browser_for_client(client_id, ws):
	if not client_id:
		return

	if client_id not in client_to_browsers:
		client_to_browsers[client_id] = set()

	client_to_browsers[client_id].add(ws)


def remove_browser_ws(ws):
	stale_ids = []

	for cid, browser_set in list(client_to_browsers.items()):
		if ws in browser_set:
			browser_set.discard(ws)

		if not browser_set:
			stale_ids.append(cid)

	for cid in stale_ids:
		client_to_browsers.pop(cid, None)


async def cleanup_loop():
	while True:
		await asyncio.sleep(60)

		for cid, ws in list(clients.items()):
			if getattr(ws, "closed", False):
				clients.pop(cid, None)
				client_to_browsers.pop(cid, None)

		for cid, browser_set in list(client_to_browsers.items()):
			alive = set()

			for browser_ws in browser_set:
				if not getattr(browser_ws, "closed", False):
					alive.add(browser_ws)

			if alive:
				client_to_browsers[cid] = alive
			else:
				client_to_browsers.pop(cid, None)


async def handle_client(ws, msg):
	custom_client_id = (msg.get("id") or "").strip()

	if not custom_client_id:
		custom_client_id = new_client_id()

	clients[custom_client_id] = ws

	await send_json(ws, {
		"type": "hello",
		"id": custom_client_id
	})

	async for raw in ws:
		try:
			data = json.loads(raw)
		except:
			continue

		msg_type = (data.get("type") or "").strip()

		if msg_type == "keepalive":
			await send_json(ws, {
				"type": "keepalive_ack"
			})
			continue

		target_browsers = list(client_to_browsers.get(custom_client_id, set()))

		if not target_browsers:
			continue

		payload = dict(data)
		payload["id"] = custom_client_id

		for browser_ws in target_browsers:
			await send_json(browser_ws, payload)


async def handle_browser(ws):
	await send_json(ws, {
		"type": "hello",
		"message": "browser connected"
	})

	async for raw in ws:
		try:
			data = json.loads(raw)
		except:
			await send_json(ws, {
				"type": "error",
				"message": "invalid json"
			})
			continue

		if data.get("type") != "cmd":
			continue

		target_id = (data.get("id") or "").strip()
		target_ws = clients.get(target_id)

		if not target_id:
			await send_json(ws, {
				"type": "error",
				"message": "missing client id"
			})
			continue

		if not target_ws:
			await send_json(ws, {
				"type": "error",
				"message": f"client not found: {target_id}"
			})
			continue

		add_browser_for_client(target_id, ws)
		await send_json(target_ws, data)


async def handler(ws):
	role = None
	custom_client_id = None

	try:
		first = await ws.recv()

		try:
			msg = json.loads(first)
		except:
			await send_json(ws, {
				"type": "error",
				"message": "invalid json"
			})
			return

		role = (msg.get("role") or "").strip().lower()

		if role == "client":
			custom_client_id = (msg.get("id") or "").strip()
			await handle_client(ws, msg)

		elif role == "browser":
			await handle_browser(ws)

		else:
			await send_json(ws, {
				"type": "error",
				"message": "unknown role"
			})

	except websockets.ConnectionClosed:
		pass
	except Exception as e:
		print("handler error:", e)
	finally:
		if role == "client":
			if custom_client_id and clients.get(custom_client_id) is ws:
				clients.pop(custom_client_id, None)

			if custom_client_id:
				client_to_browsers.pop(custom_client_id, None)

		remove_browser_ws(ws)


async def main():
	asyncio.create_task(cleanup_loop())

	async with websockets.serve(
		handler,
		"0.0.0.0",
		8888,
		max_size=WS_MAX_SIZE,
		ping_interval=20,
		ping_timeout=20
	):
		await asyncio.Future()


if __name__ == "__main__":
	asyncio.run(main())