import json
import socket
import struct
import threading
import time

services = {
    'ssh': 22,
    'alist': 5244,
    'minecraft': 25565
}
co_server = 'www.awin-x.top'
co_port = 5000

tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
tcp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
tcp_sock.bind(('0.0.0.0', 0))
tcp_sock.listen(5)
listen_port = tcp_sock.getsockname()[1]
print(f"监听端口: {listen_port}")


def pack_data(action_, message_, data_):
    # 将字符串编码为UTF-8字节串
    action_bytes = action_.encode('utf-8')
    message_bytes = message_.encode('utf-8')
    # 使用大端序打包长度信息(4字节无符号整数)
    # >I: 大端序无符号整型(4字节)
    packed = struct.pack('>I', len(action_bytes))
    packed += action_bytes
    packed += struct.pack('>I', len(message_bytes))
    packed += message_bytes
    packed += data_  # 直接附加二进制数据
    return packed


def unpack_data(packed):
    # 解析action长度(前4字节)
    action_len = struct.unpack('>I', packed[:4])[0]
    offset = 4
    # 提取action字节并解码
    action_bytes = packed[offset:offset + action_len]
    action_ = action_bytes.decode('utf-8')
    offset += action_len
    # 解析message长度(接下来的4字节)
    message_len = struct.unpack('>I', packed[offset:offset + 4])[0]
    offset += 4
    # 提取message字节并解码
    message_bytes = packed[offset:offset + message_len]
    message_ = message_bytes.decode('utf-8')
    offset += message_len
    # 剩余部分是原始二进制数据
    data_ = packed[offset:]
    return action_, message_, data_


def send_data(conn, action_, message_, data_):
    packed = pack_data(action_, message_, data_)
    conn.sendall(packed)


def recv_data(conn):
    data_ = conn.recv(4096)
    action_, message_, data_ = unpack_data(data_)
    # print(f"收到来自 {conn.getpeername()} 的数据: {data.decode()}")
    return action_, message_, data_


def handle_hello(addr, message, data):
    pass


def handle_bye(addr, message, data):
    pass


def handle_punch_to(addr, message, data):
    pass


actions = {
    'hello': handle_hello,
    'bye': lambda: print("收到来自 {client_addr} 的断开连接请求"),
    'punch_to': lambda: print("收到来自 {client_addr} 的打洞请求"),
}


def heart_beat_thread():
    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', listen_port))
    heart_beat_pack = pack_data('heart_beat', 'json_data', json.dumps({
        'services': services,
    }).encode())
    while True:
        udp_sock.sendto(heart_beat_pack, (co_server, co_port))
        time.sleep(10)


def co_server_thread(co_server_sock, co_server_addr):
    while True:
        try:
            action, message, data = recv_data(co_server_sock)
        except Exception as e:
            print(e)
            send_data(co_server_addr, 'error', f"无法解析数据", b'')
            break
        if action in actions:
            actions[action](co_server_addr, message, data)
        else:
            send_data(co_server_addr, 'error', f"未知操作: {action}", b'')
    co_server_sock.close()


def client_thread(client_sock, client_addr):
    while True:
        try:
            action, message, data = recv_data(client_sock)
        except Exception as e:
            print(e)
            send_data(client_sock, 'error', f"无法解析数据", b'')
            break
        if action in actions:
            actions[action](client_sock, client_addr, message, data)
        else:
            send_data(client_sock, 'error', f"未知操作: {action}", b'')
    client_sock.close()


if __name__ == '__main__':
    # 心跳线程
    threading.Thread(target=heart_beat_thread, daemon=True).start()

    # 等待协服务器
    while True:
        co_server_sock, co_server_addr = tcp_sock.accept()
        action, message, data = recv_data(co_server_sock)
        if action == 'hello':
            if handle_hello(co_server_addr, message, data) == 'co_server':
                print(f'协服务器{co_server_addr}连接成功')
                threading.Thread(target=co_server_thread, daemon=True).start()
                break
        co_server_sock.close()

    # 等待客户端
    while True:
        client_sock, client_addr = tcp_sock.accept()
        action, message, data = recv_data(client_sock)
        print(f"收到来自 {client_addr} 的连接")
        if action == 'hello':
            if handle_hello(co_server_addr, message, data) == 'client':
                print(f"确认客户端 {client_addr}")
                threading.Thread(target=client_thread, daemon=True).start()
                break
        client_sock.close()