import toml
from loguru import logger
import zmq
import numpy as np
import cv2

from infer import Person_model_infer
from visualize import visualize_box_mask

labels = [
    "pedestrian"
]

if __name__ == "__main__":
    cfg = toml.load('../cfg_infer_server.toml')

    # Configure log output
    logger.add(cfg['debug']['logger_filename'], format=cfg['debug']['logger_format'], retention=5, level="INFO")

    # Initialize YOLO inference model
    predictor = Person_model_infer()
    logger.info("person yolo model load success")

    context1 = zmq.Context()
    camera1_socket = context1.socket(zmq.REQ)
    camera1_socket.connect(f"tcp://localhost:{cfg['camera']['camera1_port']}")
    logger.info("connect camera1 success")

    context2 = zmq.Context()
    socket = context2.socket(zmq.REP)
    socket.bind("tcp://*:7778")
    logger.info("bind server success")
    while True:
        message = socket.recv()
        camera1_socket.send_string("")
        message = camera1_socket.recv()
        np_array = np.frombuffer(message, dtype=np.uint8)
        frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
        result = predictor.infer(frame)
        img = visualize_box_mask(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR),result,labels)
        showim = np.array(img)
        socket.send_pyobj(showim)
        if cv2.waitKey(1) == 27:
            break
    camera1_socket.close()
    socket.close()
    context1.term()
    context2.term()

    logger.info("Interrupt received, stopping...")
    logger.info("person yolo infer server exit")