import socket
import threading
import pickle
import time
import uuid

HOST = '127.0.0.1'
PORT = 65432
MOVE_SPEED = 200  # pixels per second
TICKRATE = 64     # ticks per second
MAX_USERS = 100

class Event:
    def __init__(self, name, payload):
        self.name = name
        self.payload = payload

def recvall(sock, n):
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            return None
        data += packet
    return data

def send_event(sock, event):
    data = pickle.dumps(event)
    length = len(data)
    sock.sendall(length.to_bytes(4, byteorder='big') + data)

def recv_event(sock):
    raw_len = recvall(sock, 4)
    if not raw_len:
        return None
    msg_len = int.from_bytes(raw_len, byteorder='big')
    data = recvall(sock, msg_len)
    return pickle.loads(data)

# Global dictionaries
event_subscriptions = {}   # event name -> list of client sockets
client_ids = {}            # conn -> client_id
client_states = {}         # conn -> {"client_id", "username", "x", "y", "keys", "mouse_x", "mouse_y"}
lock = threading.Lock()

def handle_client(conn, addr):
    client_id = str(uuid.uuid4())
    with lock:
        client_ids[conn] = client_id
        # Initialize state for new client.
        client_states[conn] = {
            "client_id": client_id,
            "username": "Guest",
            "x": 100,
            "y": 100,
            "keys": set(),
            "mouse_x": 0,
            "mouse_y": 0
        }
    print(f"[NEW CONNECTION] {addr} connected as {client_id}")
    
    try:
        # Send self_id so the client knows its unique id.
        send_event(conn, Event("self_id", {"client_id": client_id}))
        # Send server info.
        server_info = {
            "tickrate": TICKRATE,
            "total_users": len(client_states),
            "max_users": MAX_USERS,
            "server_ip": HOST,
            "server_port": PORT,
            "server_name": "My Multiplayer Server"
        }
        send_event(conn, Event("server_info", server_info))
    except Exception as e:
        print(f"[ERROR] sending self_id/server_info: {e}")
    
    # Broadcast that a new client connected.
    broadcast_event("client_connect", {"client_id": client_id})
    
    try:
        while True:
            event = recv_event(conn)
            if event is None:
                break
            if event.name == "register_event":
                event_name = event.payload.get('event')
                with lock:
                    event_subscriptions.setdefault(event_name, []).append(conn)
                print(f"[REGISTER] {client_id} subscribed to '{event_name}'")
            elif event.name == "set_username":
                username = event.payload.get('username', 'Guest')
                with lock:
                    if conn in client_states:
                        client_states[conn]['username'] = username
                print(f"[USERNAME] {client_id} set username to '{username}'")
            elif event.name == "keydown":
                key = event.payload.get('key')
                with lock:
                    if conn in client_states:
                        client_states[conn]['keys'].add(key)
                print(f"[KEYDOWN] {client_id} key '{key}' pressed")
            elif event.name == "keyup":
                key = event.payload.get('key')
                with lock:
                    if conn in client_states and key in client_states[conn]['keys']:
                        client_states[conn]['keys'].remove(key)
                print(f"[KEYUP] {client_id} key '{key}' released")
            elif event.name == "mouse_move":
                x = event.payload.get('x', 0)
                y = event.payload.get('y', 0)
                with lock:
                    if conn in client_states:
                        client_states[conn]['mouse_x'] = x
                        client_states[conn]['mouse_y'] = y
                # Optionally log mouse moves.
                # print(f"[MOUSE_MOVE] {client_id} moved mouse to ({x}, {y})")
            elif event.name == "ping":
                timestamp = event.payload.get('timestamp')
                send_event(conn, Event("pong", {"timestamp": timestamp}))
            # You can handle additional events here.
    except Exception as e:
        print(f"[ERROR] Exception with {client_id}: {e}")
    finally:
        with lock:
            for subs in event_subscriptions.values():
                if conn in subs:
                    subs.remove(conn)
            client_ids.pop(conn, None)
            client_states.pop(conn, None)
        print(f"[DISCONNECT] {client_id} disconnected.")
        broadcast_event("client_disconnect", {"client_id": client_id})
        conn.close()

def broadcast_event(event_name, payload):
    event = Event(event_name, payload)
    with lock:
        receivers = event_subscriptions.get(event_name, [])
        for client in receivers[:]:
            try:
                send_event(client, event)
            except Exception as e:
                print(f"[ERROR] Failed to send to client, removing: {e}")
                receivers.remove(client)

def game_loop():
    dt = 1 / TICKRATE  # Tick interval
    while True:
        time.sleep(dt)
        with lock:
            # Update each client's position based on pressed keys.
            for state in client_states.values():
                dx, dy = 0, 0
                if 'left' in state['keys']:
                    dx -= MOVE_SPEED * dt
                if 'right' in state['keys']:
                    dx += MOVE_SPEED * dt
                if 'up' in state['keys']:
                    dy -= MOVE_SPEED * dt
                if 'down' in state['keys']:
                    dy += MOVE_SPEED * dt
                state['x'] += dx
                state['y'] += dy
            # Build payload for state_update.
            players_payload = {}
            for state in client_states.values():
                players_payload[state['client_id']] = {
                    "username": state['username'],
                    "x": state['x'],
                    "y": state['y'],
                    "mouse_x": state.get('mouse_x', 0),
                    "mouse_y": state.get('mouse_y', 0)
                }
            total_users = len(client_states)
            payload = {
                "players": players_payload,
                "total_users": total_users,
                "max_users": MAX_USERS
            }
        broadcast_event("state_update", payload)

def start_server():
    print("[STARTING] Server is starting...")
    threading.Thread(target=game_loop, daemon=True).start()
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((HOST, PORT))
        s.listen()
        print(f"[LISTENING] Server is listening on {HOST}:{PORT}")
        while True:
            conn, addr = s.accept()
            threading.Thread(target=handle_client, args=(conn, addr), daemon=True).start()

if __name__ == "__main__":
    start_server()