import socket
import threading
import time
import os
import sys

# Add current directory to path so we can import ujk_vps_relay
sys.path.append(os.path.abspath(r"b:\ARCADE\ChaseChaseJokers\ujk-server"))

def test_flow():
    relay_ip = "127.0.0.1"
    
    # --- PHASE 1: START RELAY ---
    from ujk_vps_relay import RelayServer
    server = RelayServer()
    # Bind to a random port for testing
    server.main_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server.main_sock.bind(("127.0.0.1", 0))
    relay_port = server.main_sock.getsockname()[1]
    
    print(f"[*] Starting local relay on {relay_port}...")
    threading.Thread(target=server.start, daemon=True).start()
    
    time.sleep(1)
    
    # --- PHASE 2: HOST REGISTRATION ---
    host_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    host_sock.bind(("127.0.0.1", 0))
    host_addr = host_sock.getsockname()
    
    print(f"[*] Simulating Host on {host_addr}...")
    host_sock.sendto(b"HOST_HELLO", (relay_ip, relay_port))
    
    time.sleep(1)
    
    # --- PHASE 3: CLIENT CONNECT ---
    client_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client_sock.bind(("127.0.0.1", 0))
    client_addr = client_sock.getsockname()
    
    print(f"[*] Simulating Client on {client_addr}...")
    client_sock.sendto(b"HELLO_FROM_CLIENT", (relay_ip, relay_port))
    
    # --- PHASE 4: VERIFY HOST RECEIVES ---
    print("[*] Waiting for Host to receive...")
    host_sock.settimeout(5)
    try:
        data, addr = host_sock.recvfrom(65535)
        print(f"[+] Host received: '{data.decode()}' from {addr}")
        
        if data == b"HELLO_FROM_CLIENT":
            print("[SUCCESS] Client -> Relay -> Host works!")
        else:
            print(f"[FAIL] Host received unexpected data: {data}")
            return
    except socket.timeout:
        print("[FAIL] Host timed out waiting for client packet.")
        return

    # --- PHASE 5: HOST RESPONDS ---
    print(f"[*] Host replying to {addr} (The mapping port)...")
    host_sock.sendto(b"REPLY_FROM_HOST", addr)
    
    # --- PHASE 6: VERIFY CLIENT RECEIVES ---
    print("[*] Waiting for Client to receive...")
    client_sock.settimeout(5)
    try:
        data, addr = client_sock.recvfrom(65535)
        print(f"[+] Client received: '{data.decode()}' from {addr}")
        
        if data == b"REPLY_FROM_HOST":
            print("[SUCCESS] Host -> Relay -> Client works!")
        else:
            print(f"[FAIL] Client received unexpected data: {data}")
    except socket.timeout:
        print("[FAIL] Client timed out waiting for host reply.")

if __name__ == "__main__":
    test_flow()