import toml
import threading
from loguru import logger
import logging
import zmq
import numpy as np
import cv2
import time

# Custom imports
from infer_new import Yolo_model_infer
from visualize import visualize_box_mask
# Initialize locks
lock1 = threading.Lock()
lock2 = threading.Lock()
lock3 = threading.Lock()

# Global variables
src_camera_id = 1
response = {'code': 0, 'data': []}
frame = None
start = False
exit_event = threading.Event()





labels = [
    "tplatform", "tower", "sign", "shelter", "hospital", "basket", "base", 
    "Yball", "Spiller", "Rmark", "Rblock", "Rball", "Mpiller", 
    "Lpiller", "Lmark", "Bblock", "Bball"
]

# Handle server response data
def server_resp(yolo_infer_port):
    logger.info("yolo server thread init success")
    global response
    global src_camera_id

    context = zmq.Context()
    # Start server
    socket = context.socket(zmq.REP)
    socket.bind(f"tcp://*:{yolo_infer_port}")
    logger.info("yolo infer server init success")
    while not exit_event.is_set():
        try:
            message = socket.recv_string()
            # Send character 1 and 2 to switch camera, empty string requests inference data
            if message != '':
                with lock1:
                    logger.error(message)
                    src_camera_id = int(message)
                    logger.info("switch camera")
                    socket.send_pyobj(response) 
            else:
                with lock2:
                    socket.send_pyobj(response)
                    response['data'] = np.array([])
        except zmq.Again:
            time.sleep(0.01)

    socket.close()
    context.term()

# Handle camera data
def camera_resp(camera1_port, camera2_port):
    global frame
    global src_camera_id
    global start
    context = zmq.Context()
    camera1_socket = context.socket(zmq.REQ)
    camera1_socket.connect(f"tcp://localhost:{camera1_port}")
    logger.info("connect camera1 success")
    context1 = zmq.Context()
    camera2_socket = context1.socket(zmq.REQ)
    camera2_socket.connect(f"tcp://localhost:{camera2_port}")
    logger.info("connect camera2 success")

    while not exit_event.is_set():
        with lock1:
            try:
                if src_camera_id == 1:
                    camera1_socket.send_string("")
                    message = camera1_socket.recv()
                else:
                    camera2_socket.send_string("")
                    message = camera2_socket.recv()
                np_array = np.frombuffer(message, dtype=np.uint8)
                with lock3:
                    frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
                    start = True
            except:
                time.sleep(0.01)

    camera1_socket.close()
    camera2_socket.close()
    context.term()
    context1.term()


def yolo_infer_server_main(queue):
    if queue != None:
        class Handler(logging.Handler):
            def emit(self, record):
                log_entry = self.format(record)
                queue.put({'level': record.levelname.lower(), 'content': log_entry})
        # logger.remove()
        handler = Handler()
        logger.add(handler, format="{time:MM-DD HH:mm:ss} {message}", level="DEBUG")
    
    context2 = zmq.Context()
    socket_server = context2.socket(zmq.PUB)
    socket_server.setsockopt(zmq.SNDHWM,10)
    socket_server.bind("tcp://*:7777")


    cfg = toml.load('/home/evan/Workplace/project_infer/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 = Yolo_model_infer()
    logger.info("yolo model load success")

    # Start threads
    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']['front_camera_port']), daemon=True)
    mythread1.start()
    mythread2.start()
    while not exit_event.is_set():
        with lock3:
            if start:
                result = predictor.infer(frame)
                img = visualize_box_mask(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR),result,labels)
                showim = np.array(img)
                socket_server.send_pyobj(showim)
                with lock2:
                    response['data'] = result
        # time.sleep(0.01)
        if cv2.waitKey(1) == 27:
            break
    logger.info("Interrupt received, stopping...")
    exit_event.set()
    mythread1.join()
    mythread2.join()
    logger.info("yolo infer server exit")
    # try:
    #     while not exit_event.is_set():
    #         with lock3:
    #             if start:
    #                 result = predictor.infer(frame)
    #                 img = visualize_box_mask(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR),result,labels)
    #                 showim = np.array(img)
    #                 socket_server.send_pyobj(showim)
    #                 with lock2:
    #                     response['data'] = result
    #         time.sleep(0.01)
    # except KeyboardInterrupt:
    #     logger.info("Interrupt received, stopping...")
    #     exit_event.set()
    #     mythread1.join()
    #     mythread2.join()
    #     logger.info("yolo infer server exit")
if __name__ == "__main__":
    yolo_infer_server_main(None)