import toml
import threading
from loguru import logger
import zmq
from infer import Lane_model_infer
import numpy as np
import cv2
lock = threading.Lock()
response = {'code': 0, 'data': 0}

# 处理server响应数据
def server_resp(lane_infer_port):
    logger.info("lane server thread init success")
    global response
    
    context = zmq.Context()
    # 启动server
    socket = context.socket(zmq.REP)
    socket.bind(f"tcp://*:{lane_infer_port}")
    logger.info("lane infer server init success")
    while True:
        message = socket.recv_string()
        with lock:
            socket.send_pyobj(response)



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")

    # 连接摄像头server 巡线只需要连接前摄像头
    context = zmq.Context()
    camera_socket = context.socket(zmq.SUB)
    camera_socket.connect(f"tcp://localhost:{cfg['camera']['front_camera_port']}")
    camera_socket.setsockopt_string(zmq.SUBSCRIBE, "")
    logger.info("connect camera success")
    # 初始化paddle推理器
    predictor = Lane_model_infer()
    logger.info("lane model load success")
    # 启动lane_infer_server线程
    mythread = threading.Thread(target=server_resp,
                                args=(cfg['server']['lane_infer_port'],),
                                daemon=True)
    mythread.start()

    while True:
        message = camera_socket.recv()
        np_array = np.frombuffer(message, dtype=np.uint8)
        frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
        frame = cv2.resize(frame,(320,240))
        result = predictor.infer(frame)
        with lock:
            response['data'] = result
    mythread.join()
    logger.info("lane infer server exit")