import toml
import threading
from loguru import logger
import zmq
from infer_new import Yolo_model_infer
import numpy as np
import cv2
import time
lock1 = threading.Lock()
lock2 = threading.Lock()
lock3 = threading.Lock()
src_camera_id = 1
response = {'code': 0, 'data': []}
frame = None
start = False
# 处理 server 响应数据
def server_resp(yolo_infer_port):
    logger.info("yolo server thread init success")
    global response
    global src_camera_id
    
    context = zmq.Context()
    # 启动 server
    socket = context.socket(zmq.REP)
    socket.bind(f"tcp://*:{yolo_infer_port}")
    logger.info("yolo infer server init success")
    while True:
        message = socket.recv_string()
        # 发送字符 1 和 2 切换摄像头 空字符表示请求推理数据
        if message != '':
            with lock1:
                src_camera_id = int(message)
            socket.send_pyobj({'code': 0, 'data': []}) 
        else:
            with lock2:
                socket.send_pyobj(response)

# 处理摄像头数据
def camera_resp(camera1_port, camera2_port):
    global frame
    global src_camera_id
    global start
    context = zmq.Context()
    camera1_socket = context.socket(zmq.SUB)
    camera1_socket.connect(f"tcp://localhost:{camera1_port}")
    camera1_socket.setsockopt_string(zmq.SUBSCRIBE, "")
    logger.info("connect camera1 success")
    context1 = zmq.Context()
    camera2_socket = context1.socket(zmq.SUB)
    camera2_socket.connect(f"tcp://localhost:{camera2_port}")
    camera2_socket.setsockopt_string(zmq.SUBSCRIBE, "")
    logger.info("connect camera2 success")
    
    while True:
        logger.info('111')
        with lock1:
            if src_camera_id == 1:
                message = camera1_socket.recv()
            else:
                message = camera2_socket.recv()
            
            # logger.info('111')
            np_array = np.frombuffer(message, dtype=np.uint8)
            with lock3:
                frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                start = True

if __name__ == "__main__":

    cfg = toml.load('../cfg_infer_server.toml')

    # 配置日志输出
    logger.add(cfg['debug']['logger_filename'], format=cfg['debug']['logger_format'], retention = 5, level="INFO")


    # 初始化 paddle 推理器
    predictor = Yolo_model_infer()
    logger.info("yolo model load success")
    # 启动 线程 1
    mythread1 = threading.Thread(target=server_resp,
                                args=(cfg['server']['yolo_infer_port'],),
                                daemon=True)
    mythread2 = threading.Thread(target=camera_resp,
                                args=(
                                    cfg['camera']['camera1_port'],
                                    cfg['camera']['camera2_port']
                                ),
                                daemon=True)
    mythread1.start()
    mythread2.start()
    while True:
        with lock3:
            if start:
                result = predictor.infer(frame)
                with lock2:
                    response['data'] = result
    mythread1.join()
    mythread2.join()
    logger.info("yolo infer server exit")