import socket
import threading

HOST = "127.0.0.1"
PORT = 9000

devices = {}  # device_id -> socket

def handle_client(conn):
    while True:
        data = conn.recv(1024)
        if not data:
            break

        # 前 4 字节当 device_id（小端）
        device_id = int.from_bytes(data[4:8], "little")
        print("[SERVER] recv for device:", device_id, data.hex(" "))

        if device_id in devices:
            devices[device_id].send(data)
            print("[SERVER] forwarded to device", device_id)

def accept_loop():
    s = socket.socket()
    s.bind((HOST, PORT))
    s.listen()
    print("[SERVER] listening")

    while True:
        conn, _ = s.accept()

        # 第一次 recv 认为是设备注册
        first = conn.recv(1024)
        if first.startswith(b"DEVICE"):
            device_id = int(first.split(b":")[1])
            devices[device_id] = conn
            print(f"[SERVER] device {device_id} registered")
        else:
            threading.Thread(target=handle_client, args=(conn,), daemon=True).start()

accept_loop()