import socket
import threading
import pickle
import pygame
import sys
import time
import math

HOST = '127.0.0.1'
PORT = 65432

# Racing physics constants.
ACCELERATION = 300.0        # pixels per second^2
BRAKE_DECELERATION = 400.0  # pixels per second^2
FRICTION = 200.0            # pixels per second^2
TURN_SPEED = math.radians(120)  # radians per second (e.g., 120° per second)
MAX_SPEED = 400.0           # maximum forward speed
MIN_SPEED = -200.0          # maximum reverse speed

# Dimensions for the car.
CAR_WIDTH = 40
CAR_HEIGHT = 20

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 state.
players = {}         # client_id -> {"username", "x", "y", "angle", "speed"}
my_client_id = None  # Our unique client id from the server.
# Predicted state for our own car.
predicted_state = {
    "x": 100, "y": 100,
    "angle": 0,        # Facing right initially.
    "speed": 0,
    "keys": set(),
    "username": "Guest"
}
ping = 0             # measured ping in seconds
server_details = {}  # Updated from server_info and state_update events.
map_obstacles = []   # Map obstacles from server info.

def listen_to_server(sock):
    global my_client_id, players, predicted_state, ping, server_details, map_obstacles
    while True:
        try:
            event = recv_event(sock)
            if event is None:
                break
            if event.name == "self_id":
                if my_client_id is None:
                    my_client_id = event.payload.get("client_id")
                    print(f"[INFO] My client ID: {my_client_id}")
            elif event.name == "server_info":
                server_details = event.payload
                map_obstacles = server_details.get("map", [])
                print(f"[INFO] Received server info: {server_details}")
            elif event.name == "client_connect":
                print(f"[INFO] Client connected: {event.payload.get('client_id')}")
            elif event.name == "client_disconnect":
                cid = event.payload.get("client_id")
                print(f"[INFO] Client disconnected: {cid}")
                if cid in players:
                    del players[cid]
            elif event.name == "state_update":
                # Expect players to have "x", "y", "angle", "speed", "username".
                players = event.payload.get("players", {})
                # Update dynamic server details.
                if "total_users" in event.payload:
                    server_details["total_users"] = event.payload["total_users"]
                if "max_users" in event.payload:
                    server_details["max_users"] = event.payload["max_users"]
                if my_client_id and my_client_id in players:
                    server_x = players[my_client_id].get("x", 100)
                    server_y = players[my_client_id].get("y", 100)
                    server_angle = players[my_client_id].get("angle", 0)
                    server_speed = players[my_client_id].get("speed", 0)
                    tickrate = server_details.get("tickrate", 64)
                    correction_factor = 0.1 * (64 / tickrate)
                    if not predicted_state["keys"]:
                        predicted_state["x"] = server_x
                        predicted_state["y"] = server_y
                        predicted_state["angle"] = server_angle
                        predicted_state["speed"] = server_speed
                    else:
                        predicted_state["x"] += (server_x - predicted_state["x"]) * correction_factor
                        predicted_state["y"] += (server_y - predicted_state["y"]) * correction_factor
                        predicted_state["angle"] += (server_angle - predicted_state["angle"]) * correction_factor
                        predicted_state["speed"] += (server_speed - predicted_state["speed"]) * correction_factor
            elif event.name == "pong":
                sent_time = event.payload.get("timestamp")
                if sent_time:
                    round_trip = time.time() - sent_time
                    ping = round_trip
                    print(f"[PING] {ping*1000:.0f} ms")
        except Exception as e:
            print(f"[ERROR] {e}")
            break

def send_message(sock, event):
    send_event(sock, event)

def ping_loop(sock):
    while True:
        ts = time.time()
        send_message(sock, Event("ping", {"timestamp": ts}))
        time.sleep(1)

def draw_car(surface, x, y, angle, color):
    # Define a simple rectangle for the car.
    half_w = CAR_WIDTH / 2
    half_h = CAR_HEIGHT / 2
    # Corners of the rectangle before rotation.
    corners = [
        (-half_w, -half_h),
        (half_w, -half_h),
        (half_w, half_h),
        (-half_w, half_h)
    ]
    rotated = []
    cos_a = math.cos(angle)
    sin_a = math.sin(angle)
    for cx, cy in corners:
        rx = cx * cos_a - cy * sin_a
        ry = cx * sin_a + cy * cos_a
        rotated.append((int(x + rx), int(y + ry)))
    pygame.draw.polygon(surface, color, rotated)

def main():
    global my_client_id, predicted_state
    username = input("Enter your username: ")
    predicted_state["username"] = username

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((HOST, PORT))

    threading.Thread(target=listen_to_server, args=(sock,), daemon=True).start()
    threading.Thread(target=ping_loop, args=(sock,), daemon=True).start()

    # Register for events.
    for event_name in ["client_connect", "client_disconnect", "state_update"]:
        send_message(sock, Event("register_event", {"event": event_name}))
        print(f"[INFO] Registered for event: {event_name}")

    send_message(sock, Event("set_username", {"username": username}))

    pygame.init()
    screen = pygame.display.set_mode((800, 600))
    pygame.display.set_caption("Top Down Racing Multiplayer")
    clock = pygame.time.Clock()

    local_keys = set()

    while True:
        dt = clock.tick(60) / 1000.0  # Delta time in seconds

        for ev in pygame.event.get():
            if ev.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif ev.type == pygame.KEYDOWN:
                key_name = None
                if ev.key == pygame.K_LEFT:
                    key_name = "left"
                elif ev.key == pygame.K_RIGHT:
                    key_name = "right"
                elif ev.key == pygame.K_UP:
                    key_name = "up"
                elif ev.key == pygame.K_DOWN:
                    key_name = "down"
                if key_name and key_name not in local_keys:
                    local_keys.add(key_name)
                    predicted_state["keys"].add(key_name)
                    send_message(sock, Event("keydown", {"key": key_name}))
            elif ev.type == pygame.KEYUP:
                key_name = None
                if ev.key == pygame.K_LEFT:
                    key_name = "left"
                elif ev.key == pygame.K_RIGHT:
                    key_name = "right"
                elif ev.key == pygame.K_UP:
                    key_name = "up"
                elif ev.key == pygame.K_DOWN:
                    key_name = "down"
                if key_name and key_name in local_keys:
                    local_keys.remove(key_name)
                    if key_name in predicted_state["keys"]:
                        predicted_state["keys"].remove(key_name)
                    send_message(sock, Event("keyup", {"key": key_name}))

        # Racing physics update.
        # Acceleration/Braking.
        if "up" in predicted_state["keys"]:
            predicted_state["speed"] += ACCELERATION * dt
        elif "down" in predicted_state["keys"]:
            predicted_state["speed"] -= BRAKE_DECELERATION * dt
        else:
            # Apply friction to gradually slow the car.
            if predicted_state["speed"] > 0:
                predicted_state["speed"] -= FRICTION * dt
                if predicted_state["speed"] < 0:
                    predicted_state["speed"] = 0
            elif predicted_state["speed"] < 0:
                predicted_state["speed"] += FRICTION * dt
                if predicted_state["speed"] > 0:
                    predicted_state["speed"] = 0

        # Clamp speed.
        if predicted_state["speed"] > MAX_SPEED:
            predicted_state["speed"] = MAX_SPEED
        if predicted_state["speed"] < MIN_SPEED:
            predicted_state["speed"] = MIN_SPEED

        # Steering.
        if "left" in predicted_state["keys"]:
            predicted_state["angle"] -= TURN_SPEED * dt
        if "right" in predicted_state["keys"]:
            predicted_state["angle"] += TURN_SPEED * dt

        # Update position based on current speed and angle.
        predicted_state["x"] += predicted_state["speed"] * math.cos(predicted_state["angle"]) * dt
        predicted_state["y"] += predicted_state["speed"] * math.sin(predicted_state["angle"]) * dt

        # Camera follows the local car.
        local_x = predicted_state["x"]
        local_y = predicted_state["y"]
        screen_width, screen_height = 800, 600
        camera_offset = (local_x - screen_width // 2, local_y - screen_height // 2)

        # Drawing.
        screen.fill((30, 30, 30))  # Background.

        # Draw map obstacles (if any).
        for obs in map_obstacles:
            rect = pygame.Rect(obs["x"] - camera_offset[0], obs["y"] - camera_offset[1],
                               obs["width"], obs["height"])
            pygame.draw.rect(screen, (100, 100, 100), rect)

        # Draw players as cars.
        for cid, data in players.items():
            # Get car position and orientation.
            px = data.get("x", 100)
            py = data.get("y", 100)
            angle = data.get("angle", 0)
            uname = data.get("username", "Guest")
            # For our own car, use the predicted state.
            if cid == my_client_id:
                px = predicted_state["x"]
                py = predicted_state["y"]
                angle = predicted_state["angle"]
                color = (0, 0, 255)  # Blue for local player.
            else:
                color = (255, 0, 0)  # Red for others.
            draw_x = px - camera_offset[0]
            draw_y = py - camera_offset[1]
            draw_car(screen, draw_x, draw_y, angle, color)
            # Draw username above the car.
            font = pygame.font.SysFont(None, 20)
            name_surface = font.render(uname, True, (255, 255, 255))
            screen.blit(name_surface, (draw_x - name_surface.get_width() // 2, draw_y - 30))

        # Display ping and HUD info.
        ping_font = pygame.font.SysFont(None, 24)
        ping_surface = ping_font.render(f"Ping: {int(ping * 1000)} ms", True, (255, 255, 0))
        screen.blit(ping_surface, (10, 10))
        
        hud_font = pygame.font.SysFont(None, 24)
        total_users = server_details.get("total_users", len(players))
        max_users = server_details.get("max_users", "?")
        tickrate = server_details.get("tickrate", 64)
        hud_text = hud_font.render(
            f"Speed: {int(predicted_state['speed'])} px/s | Users: {total_users}/{max_users} | Client ID: {my_client_id or 'N/A'} | Tickrate: {tickrate} tps",
            True,
            (255, 255, 255)
        )
        screen.blit(hud_text, (10, 40))
        
        pygame.display.flip()

if __name__ == "__main__":
    main()