import socket
import struct
import json
import sys

#SERVER_ID = sys.argv[1]  # "A" or "B"
SERVER_ID = "A"

MCAST_GRP = "239.0.0.222"
MCAST_PORT = 50000

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("", MCAST_PORT))

mreq = struct.pack("4sl", socket.inet_aton(MCAST_GRP), socket.INADDR_ANY)
sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)

print(f"Python server {SERVER_ID} listening...")

while True:
    data, addr = sock.recvfrom(2048)
    msg = json.loads(data.decode())

    # Ignore messages not intended for this server
    if msg["dest"] not in (SERVER_ID, "ALL"):
        continue

    print(f"{SERVER_ID} received from {addr}: {msg}")

    # Process request
    if msg["cmd"] == "GO":
        result = "012345"
    elif msg["cmd"] == "RESULT":
        result = f"ACK:{msg['data']}"
    else:
        result = "UNKNOWN_CMD"

    response = {
        "type": "response",
        "src": SERVER_ID,
        "cmd": msg["cmd"],
        "data": result
    }

    sock.sendto(json.dumps(response).encode(), addr)
    print(f"{SERVER_ID} sent response to {addr}")