import threading
import time
import json
import zmq


class zmqClient:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.client_socket = None
        self.running = False
        self.zmq_server = None  # Reference to zmqServer for Unity communication
        # 记录客户端连接前的状态
        self.state = {
            'status_code': None,
            'energy': None
        }

    def set_zmq_server(self, server):
        """Set the zmqServer instance to forward messages to Unity"""
        self.zmq_server = server

    def connect(self):
        # 创建 ZeroMQ 上下文
        self.context = zmq.Context()
        # 创建 REQ 套接字（请求端）
        self.client_socket = self.context.socket(zmq.DEALER)
        # client_id = b'client1'
        # self.client_socket.setsockopt(zmq.IDENTITY,client_id)
        self.client_socket.connect(f"tcp://{self.host}:{self.port}")  # 连接到服务器
        self.running = True

    def send_to_all(self, method,params):
        if method in self.state.keys():
            self.state[method] = params
            
        # Also send to Unity via zmqServer if connected
        if self.zmq_server:
            self.zmq_server.broadcast_message(method, params)
            
        try:
            if self.running and self.client_socket != None:
                msg = {'method': method, 'params': params}
                if method in ['single_trial_plot', 'miReport']:
                    print(f"{{'method': '{method}', 'params': <Base64 Image Data>}}")
                else:
                    print(msg)
                self.client_socket.send_multipart([b'', json.dumps(msg).encode('utf-8')])
            else:
                if method in self.state.keys():
                    self.state[method] = params
        except ConnectionResetError:
            print("Connection lost.")
            self.running = False
        except Exception as e:
            print(f"An error occurred: {e}")

    def close_connection(self):
        self.running = False
        self.client_socket.close()
        self.context.term()
        print("Client closed explicitly.")
# 使用TCP客户端
if __name__ == "__main__":
    client = zmqClient('127.0.0.1', 8099)
    client.connect()
    # client.close_connection()