import pynng
import asyncio
import uuid
import time

class ImageClient:
    def __init__(self, server_address="tcp://s1.brisky.space:7701", client_id=None):
        self.server_address = server_address
        self.client_id = client_id or f"client_{uuid.uuid4().hex[:8]}"
        self.socket = None

    async def connect(self):
        """连接到服务器"""
        self.socket = pynng.Req0()
        self.socket.dial(self.server_address, block=False)
        print(f"Client {self.client_id} connected to server")

    async def send_request(self):
        """发送请求到服务器"""
        try:
            # 发送客户端 ID 作为请求
            client_id_bytes = self.client_id.encode('utf-8')
            self.socket.send(client_id_bytes)
            
            # 等待响应
            response = self.socket.recv()
            print(f"Client {self.client_id}: Received response: {response}")
            return response
        except Exception as e:
            print(f"Client {self.client_id} error: {e}")
            return None

    async def request(self):
        """发送请求并等待响应"""
        try:
            await self.connect()
            response = await self.send_request()
            return response
        except Exception as e:
            print(f"Client {self.client_id} error: {e}")

    def request_sync(self):
        """发送同步请求到服务器"""
        socket = None
        try:
            # 每次都创建一个新的 socket 实例以避免复用已关闭的 socket
            socket = pynng.Req0()
            socket.dial(self.server_address, block=True)  # 使用阻塞模式确保连接建立
            client_id_bytes = self.client_id.encode('utf-8')
            socket.send(client_id_bytes)
            response = socket.recv()
            return response
        except Exception as e:
            print(f"Client {self.client_id} error: {e}")
            return None
        finally:
            # 确保 socket 在使用后被正确关闭
            if socket:
                try:
                    socket.close()
                except:
                    pass


    async def run(self, request_interval=0.033):  # ~30fps
        """运行客户端请求循环"""
        await self.connect()
        
        request_count = 0
        while True:
            try:
                response = await self.send_request()
                if response:
                    print(f"Client {self.client_id}: Request {request_count} completed")
                
                request_count += 1
                await asyncio.sleep(request_interval)
                if request_count >= 1:
                    break
                
            except KeyboardInterrupt:
                print(f"Client {self.client_id} interrupted by user")
                break
            except Exception as e:
                print(f"Client {self.client_id} error in main loop: {e}")
                await asyncio.sleep(0.02)

async def main():
    # 创建多个客户端实例
    num_clients = 1  # 可以根据需要调整客户端数量
    tasks = []
    
    for i in range(num_clients):
        client = ImageClient(client_id=f"client_{i}_{uuid.uuid4().hex[:8]}")
        task = asyncio.create_task(client.run())
        tasks.append(task)
    
    # 等待所有客户端任务完成
    await asyncio.gather(*tasks, return_exceptions=True)

if __name__ == "__main__":
    asyncio.run(main())