import socket
import threading
import uuid

from utils import *

co_server_port = 5000
tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
tcp_sock.bind(('0.0.0.0', co_server_port))
tcp_sock.listen(5)

servers = {}

clients = {}

co_server_token = f"co_server_{uuid.uuid4().hex[:8]}"
my_charactor = 'co_server'


def handle_punch_request(client_name, server_name):
    client = clients[client_name]
    server = servers[server_name]
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    udp_sock.bind(('0.0.0.0', 5002))
    udp_sock.settimeout(20)
    co_server_punch_port = udp_sock.getsockname()[1]
    print(f'向客户端和服务端发起端口暴露请求到{co_server_punch_port}')
    send_action_tcp(server['sock'], 'punch_request', {
        'co_server_punch_port': co_server_punch_port
    })
    send_action_tcp(client['sock'], 'punch_request', {
        'co_server_punch_port': co_server_punch_port
    })
    client_punch_port = 0
    server_punch_port = 0
    count = 0
    while client_punch_port == 0 or server_punch_port == 0:
        action, message, data, addr = recv_udp(udp_sock)
        if action == 'server_punch_port':
            print(f'获取到服务端打洞端口{addr[1]}')
            server_punch_port = addr[1]
        if action == 'client_punch_port':
            print(f'获取到客户端打洞端口{addr[1]}')
            client_punch_port = addr[1]
        time.sleep(0.1)
        count += 1
    udp_sock.close()
    print('发送打洞请求')
    send_action_tcp(client['sock'], 'punch_to', {
        'server_name': server_name,
        'server_token': server['token'],
        'target_addr': (server['addr'][0], server_punch_port)
    })
    send_action_tcp(server['sock'], 'punch_to', {
        'client_name': client_name,
        'client_token': client['token'],
        'target_addr': (client['addr'][0], client_punch_port)
    })


def handle_server(sock, addr, server_info):
    server_name = server_info['name']
    server_token = server_info['token']
    if server_name in servers:
        send_action_tcp(sock, 'bye', 'server already exists')
        return
    send_action_tcp(sock, 'co_server_hello', {
        'charactor': my_charactor,
        'token': co_server_token,
    })
    servers[server_name] = {
        'sock': sock,
        'addr': addr,
        'name': server_name,
        'token': server_token,
    }
    while True:
        try:
            action, message, data = recv_tcp(sock)
            if action == 'ping':
                send_action_tcp(sock, 'pong')
            else:
                send_action_tcp(sock, 'bye', 'unknown action')
                break
        except Exception as e:
            print(e)
            send_action_tcp(sock, 'bye', f'error: {e}')
            break
    sock.close()
    del servers[server_name]


def handle_client(sock, addr, client_info):
    client_name = client_info['name']
    client_token = client_info['token']
    clients[client_name] = {
        'sock': sock,
        'addr': addr,
        'token': client_token,
    }
    while True:
        try:
            action, message, data = recv_tcp(sock)
            if action == 'ping':
                send_action_tcp(sock, 'pong')
            elif action == 'punch_request':
                server_name = message['server_name']
                if server_name in servers:
                    handle_punch_request(client_name, server_name)
                else:
                    send_action_tcp(sock, 'error', f"服务器不存在: {server_name}")
                    break
        except Exception as e:
            print(e)
            send_action_tcp(sock, 'bye', f'error: {e}')
            break
    sock.close()
    del clients[client_name]


def handle_connect(sock, addr):
    try:
        action, message, data = recv_tcp(sock)
    except Exception as e:
        print(e)
        send_action_tcp(sock, 'error', f"无法解析数据")
        return
    if action == 'server_hello':
        handle_server(sock, addr, message)
    elif action == 'client_hello':
        handle_client(sock, addr, message)
    else:
        print(f"未知连接：{addr} -> {action}")
        sock.close()


def main():
    print(f"{my_charactor} 启动")
    while True:
        conn, addr = tcp_sock.accept()
        print(f"{my_charactor} 收到来自 {addr} 的连接")
        threading.Thread(
            target=handle_connect,
            args=(conn, addr),
            daemon=True
        ).start()


if __name__ == '__main__':
    main()