
__author__    = "Francis Liu, Matteo Turilli, Andre Merzky"
__copyright__ = "Copyright 2013-2014, The AIMES Project"
__license__   = "MIT"


import radical.utils as ru

import pprint
import impl

# -----------------------------------------------------------------------------
#
# The Bundle class has two different query modes: 
#
# DIRECT_QUERY: create a bundle manager, which creates Bundle agents, which 
#               query the systems
#
# DB_QUERY    : connect to a mongodb service and fetch data from there, where
#               a bundle daemon deposited them
DIRECT_QUERY = 'direct_query'
DB_QUERY     = 'db_query'

DEFAULT_MONGODB_URL = 'mongodb://ec2-184-72-89-141.compute-1.amazonaws.com:27017/bundle_v0_1/'
DEFAULT_MONGODB_URL = 'mongodb://localhost:27017/bundle_v0_1/'

# -----------------------------------------------------------------------------
#
class Bundle(object):
    """
    The main class, Bundle, accepts a config file and an 'origin' resource name.
    The latter is used for bandwidth information queries, which are always
    relative to the given origin.  The class uses the bundle implementation to
    create a Resource instance for each given resource.
    """

    # --------------------------------------------------------------------------
    #
    def __init__(self, config_file=None, origin=None, 
                 query_mode=DB_QUERY, 
                 mongodb_url=DEFAULT_MONGODB_URL):

        self.config_file = config_file
        self.origin      = origin
        self.mongodb_url = mongodb_url

        if  query_mode == DIRECT_QUERY :
            self.query_direct ()

        elif query_mode == DB_QUERY :
            self.query_db ()

        else :
            raise ValueError ("no such query mode '%s'" % query_mode)


    # --------------------------------------------------------------------------
    #
    def query_direct (self) :

        self.bm = impl.BundleManager()
        self.bm.load_cluster_credentials(self.config_file)

        self.resource = list()
        self._priv    = self.bm.get_data(self.origin)

      # for cluster in self._priv['cluster_list']:
      #     print "===> %s" % cluster
      #     print
      #     print "--   config"
      #     pprint.pprint(self._priv['cluster_config'][cluster])
      #     print
      #     print "--   workload"
      #     pprint.pprint(self._priv['cluster_workload'][cluster])
      #     print
      #     print "--   bandwidths"
      #     pprint.pprint(self._priv['cluster_bandwidth'][cluster])
      #     print
      #     print

        # we have a dictionary of Resources instances, indexed by resource name
        self.resources = dict()
        for resource_name in self._priv['cluster_list']:

            config     = self._priv['cluster_config'   ][resource_name]
            workload   = self._priv['cluster_workload' ][resource_name]
            bandwidths = self._priv['cluster_bandwidth'][resource_name]

            # import pprint
            # pprint.pprint(bandwidths)

            self.resources[resource_name] = Resource(resource_name, config, workload, bandwidths)


        # and a list of Queue instances, for all queues of all resources
        self.queues = list()
        for resource in self.resources:
            self.queues += self.resources[resource].queues.values()


    # --------------------------------------------------------------------------
    #
    def query_db (self) :

        mongo, db, dbname, cname, pname = ru.mongodb_connect (self.mongodb_url)

        self._priv = dict()
        self._priv['cluster_list']      = list()
        self._priv['cluster_config']    = dict()
        self._priv['cluster_workload']  = dict()
        self._priv['cluster_bandwidth'] = dict()


        for doc in list(db['config'].find ()):
            self._priv['cluster_list'].append (doc['_id'])
            self._priv['cluster_config'][doc['_id']] = doc