import socket
import json
import time
from collections import defaultdict


class CoordinatorServer:
    def __init__(self, host='0.0.0.0', port=5000):
        self.host = host
        self.port = port
        self.services = {}  # 服务名 -> (公网地址, 内网端口)
        self.clients = defaultdict(dict)  # 客户端ID -> 信息
        self.udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp_sock.bind((host, port))
        self.running = True
        print(f"协调服务器运行在 {host}:{port}")

    def handle_register(self, data, addr):
        """处理服务注册请求"""
        try:
            service_name = data['service_name']
            internal_port = data['internal_port']
            client_id = data['client_id']

            # 记录客户端信息
            self.clients[client_id] = {
                'addr': addr,
                'service_name': service_name,
                'internal_port': internal_port,
                'last_seen': time.time()
            }

            # 注册服务
            self.services[service_name] = (addr, internal_port)
            print(f"服务注册: {service_name} (端口:{internal_port}) 来自 {addr}")

            return {'status': 'success', 'message': '服务注册成功'}
        except Exception as e:
            return {'status': 'error', 'message': str(e)}

    def handle_request(self, data, addr):
        """处理服务请求"""
        try:
            service_name = data['service_name']
            client_id = data['client_id']

            if service_name not in self.services:
                return {'status': 'error', 'message': '服务未找到'}

            # 记录请求客户端信息
            self.clients[client_id] = {
                'addr': addr,
                'service_name': service_name,
                'last_seen': time.time()
            }

            # 获取服务提供者的信息
            provider_addr, internal_port = self.services[service_name]
            target_id = [
                client_id
                for client_id, info in self.clients.items()
                if info['service_name'] == service_name
            ][0]

            print(f"服务请求: {service_name} 来自 {addr}, 提供者 {provider_addr}")

            return {
                'status': 'success',
                'provider_addr': provider_addr,
                'internal_port': internal_port,
                'target_id': target_id
            }
        except Exception as e:
            return {'status': 'error', 'message': str(e)}

    def handle_punch_request(self, data, addr):
        """处理打洞请求"""
        try:
            client_id = data['client_id']
            target_id = data['target_id']

            # 获取目标客户端信息
            if target_id not in self.clients:
                return {'status': 'error', 'message': '目标客户端未找到'}

            target_addr = self.clients[target_id]['addr']

            print(f"打洞请求: {addr} -> {target_addr}")

            # 通知双方对方的地址
            self.udp_sock.sendto(json.dumps({
                'action': 'punch_request',
                'client_id': client_id,
                'client_addr': addr
            }).encode(), target_addr)

            return {
                'status': 'success',
                'target_addr': target_addr
            }
        except Exception as e:
            return {'status': 'error', 'message': str(e)}

    def run(self):
        """运行协调服务器"""
        print("协调服务器已启动，等待连接...")
        while self.running:
            try:
                data, addr = self.udp_sock.recvfrom(4096)
                try:
                    message = json.loads(data.decode())
                    action = message.get('action')

                    if action == 'register':
                        response = self.handle_register(message, addr)
                    elif action == 'request':
                        response = self.handle_request(message, addr)
                    elif action == 'punch_request':
                        response = self.handle_punch_request(message, addr)
                    else:
                        response = {'status': 'error', 'message': '无效操作'}

                    self.udp_sock.sendto(json.dumps(response).encode(), addr)
                except json.JSONDecodeError:
                    self.udp_sock.sendto(json.dumps({
                        'status': 'error',
                        'message': '无效的JSON数据'
                    }).encode(), addr)
            except Exception as e:
                print(f"服务器错误: {str(e)}")

        self.udp_sock.close()


if __name__ == '__main__':
    server = CoordinatorServer()
    server.run()