import json
from loguru import logger
import yaml
from websocket_server import WebsocketServer
import os

# 设置当前工作目录
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# 配置日志
logger.add("logs.log", rotation="1 MB", level="INFO", format="{time} - {level} - {message}")

moon_addr = None
port = None

with open('moon_config.yaml', 'r') as config_file:
    config = yaml.safe_load(config_file)
    moon_addr = config.get("addr", "127.0.0.1")
    port = config.get("port", 8765)

clients = dict()  # id -> name
clients_rev = dict()  # name -> id
clients_connection = dict()  # id -> client obj

def new_client(client, server):
    clients_connection[client["id"]] = client
    logger.info(f"新客户端连接: {client['id']}")

def client_left(client, server):
    logger.info(f"客户端断开连接: {client['id']}")
    # 连接关闭时从集合中移除
    if client in clients.values():
        clients.pop(clients_rev[client['id']])
        clients_rev.pop(client['id'])

def message_received(client, server, message):
    logger.info(f"处理信息: {message}")
    try:
        msg = json.loads(message)
        clients[msg["id"]] = client['id']
        clients_rev[client['id']] = msg["id"]
        
        if msg["receiver"] == "auth":
            return
        
        if msg["receiver"] == "*":
            try:
                for i in clients_rev.values():
                    content = msg["content"]
                    content["emitter"] = msg["id"]
                    server.send_message(
                        clients_connection[i], 
                        json.dumps(content))
            except KeyError:
                logger.warning("无法发送至端点")
        
        try:
            content = msg["content"]
            content["emitter"] = msg["id"]
            server.send_message(
                clients_connection[clients[msg["receiver"]]], 
                json.dumps(content))
        except KeyError:
            logger.warning("无法发送至端点")
    except json.JSONDecodeError:
        logger.error("无法解码 JSON 消息")

def main():
    global moon_addr
    global port
    server = WebsocketServer(host=moon_addr, port=port)
    server.set_fn_new_client(new_client)
    server.set_fn_client_left(client_left)
    server.set_fn_message_received(message_received)
    logger.info(f'启动"卫星"中继服务器: ws://{moon_addr}:{port}')
    server.run_forever()

if __name__ == "__main__":
    main()
