#!/usr/bin/env python3

import sys
import rospy
import threading

from sensor_msgs.msg import Image
from lasr_vision_msgs.srv import YoloDetection, YoloDetectionRequest

if len(sys.argv) < 2:
    print('Usage: rosrun lasr_objcet_detection_yolov8 relay <source_topic> [model.pt]')
    exit()

# figure out what we are listening to
listen_topic = sys.argv[1]

# figure out what model we are using
if len(sys.argv) >= 3:
    model = sys.argv[2]
else:
    model = "yolov8n.pt"

processing = False

def detect(image):
    global processing
    processing = True
    rospy.loginfo("Received image message")

    try:
        detect_service = rospy.ServiceProxy('/yolov8/detect', YoloDetection)
        req = YoloDetectionRequest()
        req.image_raw = image
        req.dataset = model
        req.confidence = 0.25
        req.nms = 0.4
        resp = detect_service(req)
        print(resp)
    except rospy.ServiceException as e:
        rospy.logerr("Service call failed: %s" % e)
    finally:
        processing = False

def image_callback(image):
    global processing
    if processing:
        return

    t = threading.Thread(target=detect, args=(image,))
    t.start()

def listener():
    rospy.init_node('image_listener', anonymous=True)
    rospy.wait_for_service('/yolov8/detect')
    rospy.Subscriber(listen_topic, Image, image_callback)
    rospy.spin()

if __name__ == '__main__':
    listener()
