from __future__ import annotations

import json
import socket
import threading
import time
from dataclasses import dataclass
from typing import Any


HOST = "127.0.0.1"
PORT = 7777
ACCEPT_TIMEOUT_SECONDS = 0.5

clients_lock = threading.Lock()
clients: dict[socket.socket, "Client"] = {}
next_player_id = 1


@dataclass
class Client:
    connection: socket.socket
    address: tuple[str, int]
    player_id: int

    @property
    def label(self) -> str:
        return f"{self.address[0]}:{self.address[1]}"


def log(message: str) -> None:
    print(message, flush=True)


def assign_client(connection: socket.socket, address: tuple[str, int]) -> Client:
    global next_player_id

    with clients_lock:
        # The server owns player IDs so early clients do not need to coordinate
        # identity with each other or know anything about remote players yet.
        client = Client(connection=connection, address=address, player_id=next_player_id)
        next_player_id += 1
        clients[connection] = client

    return client


def remove_client(client: Client) -> bool:
    with clients_lock:
        return clients.pop(client.connection, None) is not None


def send_packet(connection: socket.socket, packet: dict[str, Any]) -> None:
    encoded = json.dumps(packet, separators=(",", ":")).encode("utf-8") + b"\n"
    connection.sendall(encoded)


def disconnect_client(client: Client) -> None:
    if not remove_client(client):
        return

    try:
        client.connection.close()
    except OSError:
        pass

    log(f"Client disconnected: {client.label} (player {client.player_id})")
    broadcast_disconnect(client)


def handle_client(connection: socket.socket, address: tuple[str, int]) -> None:
    client = assign_client(connection, address)
    log(f"Client connected: {client.label} (player {client.player_id})")

    with connection:
        buffer = ""

        try:
            send_packet(
                connection,
                {
                    "type": "welcome",
                    "playerId": client.player_id,
                    "serverTime": time.time(),
                },
            )

            while True:
                chunk = connection.recv(4096)
                if not chunk:
                    break

                buffer += chunk.decode("utf-8", errors="replace")
                while "\n" in buffer:
                    line, buffer = buffer.split("\n", 1)
                    handle_line(client, line.strip())
        except ConnectionResetError:
            log(f"Client disconnected unexpectedly: {client.label} (player {client.player_id})")
        except OSError as error:
            log(f"Client connection error: {client.label} (player {client.player_id}): {error}")
        finally:
            disconnect_client(client)


def handle_line(client: Client, line: str) -> None:
    if not line:
        return

    try:
        packet = json.loads(line)
    except json.JSONDecodeError as error:
        log(f"Invalid JSON from {client.label}: {error}: {line}")
        return

    if not isinstance(packet, dict):
        log(f"Invalid packet from {client.label}: expected JSON object: {packet}")
        return

    if packet.get("type") == "transform":
        packet["playerId"] = client.player_id
        packet["serverTime"] = time.time()

    print_packet(client, packet)

    if packet.get("type") == "transform":
        broadcast_transform(client, packet)


def broadcast_transform(sender: Client, packet: dict[str, Any]) -> None:
    with clients_lock:
        recipients = [client for client in clients.values() if client.connection != sender.connection]

    failed_recipients: list[Client] = []
    for recipient in recipients:
        try:
            # Transforms go to every other client only. The sender already knows
            # its own movement; echoing it back would create duplicate local state.
            send_packet(recipient.connection, packet)
        except OSError:
            failed_recipients.append(recipient)

    for recipient in failed_recipients:
        disconnect_client(recipient)

    log(f"Broadcast transform from player {sender.player_id} to {len(recipients) - len(failed_recipients)} other client(s)")


def broadcast_disconnect(disconnected_client: Client) -> None:
    packet = {
        "type": "disconnect",
        "playerId": disconnected_client.player_id,
        "serverTime": time.time(),
    }

    with clients_lock:
        recipients = list(clients.values())

    failed_recipients: list[Client] = []
    for recipient in recipients:
        try:
            send_packet(recipient.connection, packet)
        except OSError:
            failed_recipients.append(recipient)

    for recipient in failed_recipients:
        disconnect_client(recipient)

    log(
        f"Broadcast disconnect for player {disconnected_client.player_id} "
        f"to {len(recipients) - len(failed_recipients)} other client(s)"
    )


def print_packet(client: Client, packet: dict[str, Any]) -> None:
    if packet.get("type") != "transform":
        log(f"Packet from {client.label}: {packet}")
        return

    try:
        x = float(packet["x"])
        y = float(packet["y"])
        z = float(packet["z"])
        angle_z = float(packet["angleZ"])
    except (KeyError, TypeError, ValueError):
        log(f"Malformed transform packet from {client.label}: {packet}")
        return

    movement_type = packet.get("movementType", "normal")
    extra_fields = []
    for field_name in ("playerId", "cellId", "worldspaceId", "clientTime", "serverTime"):
        if field_name in packet:
            extra_fields.append(f"{field_name}={packet[field_name]}")

    extra_details = ""
    if extra_fields:
        extra_details = ", " + ", ".join(extra_fields)

    log(
        "Transform from "
        f"{client.label}: x={x:.2f}, y={y:.2f}, z={z:.2f}, angleZ={angle_z:.2f}, "
        f"movementType={movement_type}{extra_details}"
    )


def main() -> None:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
        if hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
            server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
        server_socket.bind((HOST, PORT))
        server_socket.listen()
        server_socket.settimeout(ACCEPT_TIMEOUT_SECONDS)

        log(f"Fallout 4 Together local test server listening on {HOST}:{PORT}")
        log("Waiting for newline-separated JSON transform packets...")

        while True:
            try:
                connection, address = server_socket.accept()
            except socket.timeout:
                # Wake periodically so Ctrl+C is handled promptly in Windows terminals.
                continue

            thread = threading.Thread(target=handle_client, args=(connection, address), daemon=True)
            thread.start()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        log("\nServer stopped.")