import socket
import threading
import time

# --- CONFIGURATION ---
HOST_PORT = 57373
CLIENT_PORT = 57374
PORT_RANGE = range(50000, 50101)
TIMEOUT = 120  # Seconds to keep a registration alive
# --- --- --- --- ---

class RelayServer:
    def __init__(self):
        # Socket for the Host (The one who creates the sala)
        self.host_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.host_sock.bind(("0.0.0.0", HOST_PORT))
        
        # Socket for the Clients (The ones who join)
        self.client_main_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.client_main_sock.bind(("0.0.0.0", CLIENT_PORT))
        
        self.host_addr = None
        self.host_last_seen = 0
        
        # client_addr -> { "port": int, "socket": socket, "last_seen": float }
        self.clients = {}
        # relay_port -> client_addr
        self.port_to_client = {}
        
        self.lock = threading.Lock()
        self.available_ports = list(PORT_RANGE)

    def get_port(self):
        if not self.available_ports:
            return None
        return self.available_ports.pop(0)

    def release_port(self, port):
        self.available_ports.append(port)
        self.available_ports.sort()

    def start(self):
        print(f"============================================")
        print(f"   CCJ AUTOMATIC ZERO-CONFIG RELAY (ARM64)  ")
        print(f"============================================")
        print(f"[*] HOST Port: {HOST_PORT} (Point your game here to CREATE)")
        print(f"[*] CLIENT Port: {CLIENT_PORT} (Point your game here to JOIN)")
        print(f"[*] Mapping range: {PORT_RANGE.start}-{PORT_RANGE.stop-1}")
        print(f"[*] No scripts needed on PCs. Just point and play.")
        print(f"============================================")
        
        threading.Thread(target=self.host_listener, daemon=True).start()
        threading.Thread(target=self.client_listener, daemon=True).start()
        threading.Thread(target=self.cleanup_loop, daemon=True).start()
        
        while True:
            time.sleep(1)

    def host_listener(self):
        """ Registers the Host when any packet arrives at HOST_PORT. """
        while True:
            try:
                data, addr = self.host_sock.recvfrom(65535)
                now = time.time()
                
                with self.lock:
                    if self.host_addr != addr:
                        print(f"[+] NEW HOST Registered: {addr}")
                    self.host_addr = addr
                    self.host_last_seen = now
                
                # Note: We don't forward packets FROM the host port here.
                # The host sends game data TO the mapping ports (5000x) that it receives packets from.
            except Exception as e:
                print(f"[!] Host listener error: {e}")

    def client_listener(self):
        """ Handles new clients connecting to the CLIENT_PORT. """
        while True:
            try:
                data, addr = self.client_main_sock.recvfrom(65535)
                now = time.time()
                
                with self.lock:
                    if not self.host_addr:
                        # No host yet? We ignore for now.
                        continue
                    
                    if addr not in self.clients:
                        # New client identified by their IP/Port
                        v_port = self.get_port()
                        if v_port is None:
                            print(f"[!] No more ports available for client {addr}")
                            continue

                        # Create a return socket for this specific client
                        client_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                        try:
                            client_sock.bind(("0.0.0.0", v_port))
                        except Exception as e:
                            print(f"[!] Failed to bind port {v_port}: {e}")
                            self.release_port(v_port)
                            continue

                        self.clients[addr] = {
                            "port": v_port,
                            "socket": client_sock,
                            "last_seen": now
                        }
                        self.port_to_client[v_port] = addr
                        
                        # Start listening for replies from the Host to this mapping port
                        threading.Thread(target=self.host_to_client_loop, args=(client_sock, v_port), daemon=True).start()
                        print(f"[+] Client {addr} mapped to VPS:{v_port} -> Host:{self.host_addr[1]}")

                    client = self.clients[addr]
                    client["last_seen"] = now
                    
                    # Forward the client data to the Host FROM the dedicated mapping port
                    client["socket"].sendto(data, self.host_addr)
            except Exception as e:
                print(f"[!] Client listener error: {e}")

    def host_to_client_loop(self, sock, port):
        """ Listens on a mapping port (5000x) for packets from the Host and sends back to Client. """
        while True:
            try:
                data, addr = sock.recvfrom(65535)
                with self.lock:
                    # Security check: must come from the registered Host IP
                    if not self.host_addr or addr[0] != self.host_addr[0]:
                        continue
                    
                    # Auto-update Host port if the game uses a different one than registration
                    if addr != self.host_addr:
                        print(f"[*] Updating Host game port: {self.host_addr[1]} -> {addr[1]}")
                        self.host_addr = addr

                    client_addr = self.port_to_client.get(port)
                    if not client_addr:
                        break

                # Send back to the Client via the main client socket
                self.client_main_sock.sendto(data, client_addr)
            except:
                break

    def cleanup_loop(self):
        """ Removes inactive connections to free up ports. """
        while True:
            time.sleep(10)
            now = time.time()
            with self.lock:
                to_delete = []
                for addr, info in self.clients.items():
                    if now - info["last_seen"] > TIMEOUT:
                        to_delete.append(addr)
                
                for addr in to_delete:
                    info = self.clients[addr]
                    port = info["port"]
                    print(f"[-] Client {addr} timed out. Releasing port {port}")
                    try:
                        info["socket"].close()
                    except:
                        pass
                    del self.port_to_client[port]
                    del self.clients[addr]
                    self.release_port(port)

                if self.host_addr and now - self.host_last_seen > TIMEOUT:
                    print(f"[!] Host {self.host_addr} timed out!")
                    self.host_addr = None

if __name__ == "__main__":
    server = RelayServer()
    server.start()