#!/usr/bin/env python3

import sys
import rospy
import threading

import cv2
import numpy as np

# wouldn't use cv_bridge usually but i'm lazy today
from cv_bridge import CvBridge

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

if len(sys.argv) < 2:
    print('Usage: rosrun lasr_vision_yolov8 construct_mask <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-seg.pt"

# setup cv bridge
bridge = CvBridge()

# debug topic
debug_publisher1 = rospy.Publisher('/yolov8/debug_mask', Image, queue_size=1)
debug_publisher2 = rospy.Publisher('/yolov8/debug_mask_all', Image, queue_size=1)

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.0
        req.nms = 0.0
        resp = detect_service(req)

        print(resp)

        # pick the first detection as an example
        if len(resp.detected_objects) > 0:
            detection = resp.detected_objects[0]

            if len(detection.xyseg) > 0:
                # unflatten the array
                contours = np.array(detection.xyseg).reshape(-1, 2)

                # draw using opencv
                img = np.zeros((image.height, image.width), dtype=np.uint8)
                cv2.fillPoly(img, pts = [contours], color = (255,255,255))
                
                # send to topic
                img_msg = bridge.cv2_to_imgmsg(img, encoding="passthrough")
                debug_publisher1.publish(img_msg)
            else:
                print('WARN: No segmentation was performed on the image!')
        else:
            print('WARN: no detections')

        # draw all of them
        if len(resp.detected_objects) > 0:
            img = np.zeros((image.height, image.width), dtype=np.uint8)
            for detection in resp.detected_objects:
                if len(detection.xyseg) > 0:
                    contours = np.array(detection.xyseg).reshape(-1, 2)
                    r,g,b = np.random.randint(0, 255, size=3)
                    cv2.fillPoly(img, pts = [contours], color = (int(r), int(g), int(b)))
            img_msg = bridge.cv2_to_imgmsg(img, encoding="passthrough")
            debug_publisher2.publish(img_msg)
        else:
            print('WARN: no detections')

    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.Subscriber(listen_topic, Image, image_callback)
    rospy.spin()

if __name__ == '__main__':
    listener()
