#!/usr/bin/env python
#
# License: BSD
#
#   https://raw.github.com/robotics-in-concert/rocon_concert/license/LICENSE
#

##############################################################################
# Imports
##############################################################################

import socket
import argparse
import threading

import rospy
import rosgraph
import unique_id
import scheduler_msgs.msg as scheduler_msgs
import rocon_console.console as console

##############################################################################
# Functions
##############################################################################


def parse_arguments():
    parser = argparse.ArgumentParser(description='Scans for and prettifies scheduler request feedback information')
    args = parser.parse_args()
    return args


def resolve_scheduler_request_topics():
    # using rostopic.find_topic() is probably easier than this
    master = rosgraph.Master(rospy.get_name())
    unused_publishers, unused_subscribers, unused_services = master.getSystemState()
    topics = set()
    published_topics = master.getPublishedTopics(subgraph='')
    for topic_info in published_topics:
        topic_name = topic_info[0]
        topic_type = topic_info[1]
        if topic_type == 'scheduler_msgs/SchedulerRequests':
            (unused_leading, unused_separator, trailing_name) = topic_name.rpartition('/')
            # really ugly rough hack to make sure we don't listen to requester publishers (only want scheduler publishers)
            if len(trailing_name) > 16:
                topics.add(topic_name)
    return topics


def pretty_print_scheduler_requests(scheduler_requests, updated_requester_id):
    '''
      @param scheduler_requests dic of hex string : scheduler_msgs.SchedulerRequests

      @param updated_requester_id : id of the requester that received a callback.
    '''
    console.pretty_println("##############################################################################", console.bold)
    console.pretty_println("#                 Scheduler Request Feedback                                 #", console.bold)
    console.pretty_println("##############################################################################", console.bold)
    for requester_id, requests in scheduler_requests.iteritems():
        if updated_requester_id == requester_id:
            console.pretty_println("Requester [%s]" % requester_id, console.bold_yellow)
        else:
            console.pretty_println("Requester [%s]" % requester_id, console.bold)
        #print(console.cyan + "  Id" + console.reset + " : " + console.yellow + "%s" % requester_id + console.reset)
        for request in requests:
            console.pretty_println("  Request", console.green)
            print(console.cyan + "    Id       " + console.reset + " : " + console.yellow + "%s" % unique_id.toHexString(request.id) + console.reset)
            print(console.cyan + "    Priority " + console.reset + " : " + console.yellow + "%s" % str(request.priority) + console.reset)
            print(console.cyan + "    Status   " + console.reset + " : " + console.yellow + "%s" % status_strings[request.status] + console.reset)
            print(console.cyan + "    Resources" + console.reset + " :"),
            prefix = ''
            if request.resources:
                for resource in request.resources:
                    print(prefix),
                    console.pretty_print(unique_id.toHexString(resource.id), console.red)
                    print("-"),
                    console.pretty_print(resource.uri, console.green)
                    print("-"),
                    console.pretty_print(resource.rapp + "\n", console.yellow)
                    prefix = prefix if prefix else "              : "
            else:
                print('-')

##############################################################################
# Class
##############################################################################


class Gatherer(object):

    def __init__(self):
        self._lock = threading.Lock()
        self._subscribers = {}
        self._refresh()
        self._timer = rospy.Timer(rospy.Duration(5.0), self._refresh)
        self._scheduler_requests = {}

    def _refresh(self, unused_event=None):
        try:
            topics = resolve_scheduler_request_topics()
            difference = lambda l1, l2: [x for x in l1 if x not in l2]
            new_topics = difference(topics, self._subscribers.keys())
            lost_topics = difference(self._subscribers.keys(), topics)
            for topic_name in new_topics:
                self._subscribers[topic_name] = rospy.Subscriber(topic_name, scheduler_msgs.SchedulerRequests, self._listener_callback)
            for topic_name in lost_topics:
                self._subscribers[topic_name].unregister()
                del self._subscribers[topic_name]
        except socket.error:
            pass  # ros is shutting down

    def _listener_callback(self, msg):
        self._lock.acquire()
        self._scheduler_requests[unique_id.toHexString(msg.requester)] = msg.requests
        pretty_print_scheduler_requests(self._scheduler_requests, unique_id.toHexString(msg.requester))
        self._lock.release()

    def shutdown(self):
        self._timer.shutdown()

##############################################################################
# Main
##############################################################################

if __name__ == '__main__':
    status_strings = ["new", "reserved", "waiting", "granted", "preempting", "canceling", "closed"]
    rospy.init_node('concert_scheduler_requests')
    args = parse_arguments()
    gatherer = Gatherer()
    rospy.spin()
    gatherer.shutdown()
