from PySide6.QtCore import QObject, Signal, Slot, QByteArray
from PySide6.QtNetwork import QUdpSocket, QHostAddress

from loguru import logger


class UdpComm(QObject):
 dataReceived = Signal(QByteArray, QHostAddress, int) # data, sender IP, port

 def __init__(self, parent=None):
 super().__init__(parent)
 self.socket = QUdpSocket(self)
 self.socket.readyRead.connect(self.on_ready_read)
 logger.info("UdpComm 初始化完成")

 def bind(self, port: int = 0, address: str = "0.0.0.0"):
 """绑定本地端口接收数据"""
 success = self.socket.bind(QHostAddress(address), port)
 if success:
 logger.info(f"绑定UDP端口成功: {address}:{self.socket.localPort()}")
 else:
 logger.error(f"绑定UDP端口失败: {address}:{port}")

 def send_data(self, data: bytes, ip: str, port: int):
 """发送数据到指定IP和端口"""
 if isinstance(data, str):
 data = data.encode("utf-8") # 支持传入字符串,自动转为字节
 sent = self.socket.writeDatagram(data, QHostAddress(ip), port)
 # logger.info(f"发送UDP数据: {sent} 字节 → {ip}:{port}")

 @Slot()
 def on_ready_read(self):
 """接收数据回调"""
 try:
 while self.socket.hasPendingDatagrams():
 datagram, host, port = self.socket.readDatagram(self.socket.pendingDatagramSize())
 logger.trace(f"收到UDP数据: {datagram} ← {host.toString()}:{port}")
 self.dataReceived.emit(datagram, host, port)
 except Exception as e:
 logger.exception(f"UDP 读取异常: {e}")

 def close(self):
 """主动关闭 socket"""
 if self.socket:
 logger.info("关闭 UDP socket")
 self.socket.close()

 def __del__(self):
 """析构时自动关闭 socket"""
 logger.info("UdpComm 析构,准备释放资源")
 self.close()