#!/usr/bin/env python3

import actionlib
import rospy

# import msgs
import aruco.msg

# imports for aruco detection
from aruco.aruco_class import ArucoDetector

# imports for debugging
import cv2

class ActionServer(ArucoDetector):
    def __init__(self):

        ### launch parameters
        try:
            DICT = str(rospy.get_param('aruco_dict'))
            MARKER_SIZE = float(rospy.get_param('marker_size'))
            CALIB_PATH = str(rospy.get_param('calib_path'))
        except Exception:
            print("Can't load Aruco parameters from ROS parameter server")
            DICT = 'DICT_ARUCO_ORIGINAL'
            MARKER_SIZE = 0.077
            CALIB_PATH = 'data/calib_params_MS_640_360.npz'

        ### aruco detector initialization
        super().__init__(DICT, MARKER_SIZE, CALIB_PATH)

        ### action server initialization
        action_name = 'detectMarker'
        action_definition = aruco.msg.detectMarkerAction
        self.feedback = aruco.msg.detectMarkerFeedback()
        self.result = aruco.msg.detectMarkerResult()
        self.server = actionlib.SimpleActionServer(
            action_name, action_definition, self.execute, False)
        self.server.start()
        self.server.register_preempt_callback(self.setPreemptState)

        ### INFO
        rospy.loginfo("Aruco server initialization done.")
        
    '''
    Execute callback - action function - goal is action ros msg goal
    '''
    def execute(self, goal):
        
        r = rospy.Rate(10)

        self.goal_id = goal.id

        ### MAIN LOOP
        while not rospy.is_shutdown() and self.server.is_active():

            # obtain image, detect markers and estimate their poses
            self.perform_detection()
            
            # a marker detected
            if self.ids is not None:

                # the wanted marker detected
                if [self.goal_id] in self.ids:
            
                    # process detections
                    self.process_detections()
                    
                    # write results
                    self.write_results()
                    self.server.set_succeeded(self.result)

                    # INFO
                    if self.DO_DEBUGING:
                        self.cv_image = cv2.aruco.drawAxis(self.cv_image, self.cam_mtx,
                            self.dist_coefs, self.rvec, self.tvec, 3 * self.axis_length)
                            
                if self.DO_DEBUGING:
                    ids_str = ''.join(str(e) for e in self.ids)
                    rospy.loginfo("Detected IDs: %s.", ids_str)
                    self.draw_detections()

            else:
                pass # TODO write unsucessful results


            if self.DO_DEBUGING:
                self.image_pub.publish(self.bridge.cv2_to_imgmsg(self.cv_image, "bgr8"))

            # TODO if something nasty happens
            if False:
                bExecute_detection = False
                ### abort action -> send result to client - problem!
                self.result.exitcode = 1
                self.server.set_succeeded(self.result)

            # rate of while loop
            r.sleep()  
    
    '''
    setPreempt - action go to preempt state, usually called when the goal is canceled
    '''
    def setPreemptState(self):
        # TODO tidy up the preempted action, then abort the action and send the result
        print('Preempt action')
        # stop subscribing images
        self.image_sub.unregister()

        self.result.exitcode = 10
        self.server.set_succeeded(self.result)

    ''' # NOT USED NOW
    webcam_callback - called when a webcam image is subscribed
    '''
    def webcam_callback(self,data):
        # obtain image
        self.cv_image = self.bridge.imgmsg_to_cv2(data)

if __name__ == '__main__':

    node_name = 'aruco_server'
    rospy.init_node(node_name)
    server = ActionServer()
    rospy.spin()
