#!/usr/bin/env python

__author__="Subhav Pradhan, Shweta Khare"

import os, sys, time
import argparse, ConfigParser
import json
from kazoo.client import KazooClient, KazooState
from kazoo.recipe.watchers import ChildrenWatch
from kazoo.protocol.states import EventType
from pymongo import MongoClient
from chariot_runtime_libs import Serialize, invoke_management_engine, get_logger

def connection_state_listener(state):
    if state == KazooState.LOST:
        logger.error ("Connection to server lost!")
        sys.exit()
    elif state == KazooState.SUSPENDED:
        logger.error ("Connection to server has been suspended!")
        sys.exit()

def membership_watch(children,event):
    global CURRENT_MEMBERS

    if event and event.type==EventType.CHILD:
        # Handle node join.
        if len(children) > len(CURRENT_MEMBERS):
            for child in children:
                if child not in CURRENT_MEMBERS:
                    CURRENT_MEMBERS.append(child)
                    logger.info ("Node: " + child + " has joined!")
                    
                    # Get node information from zookeeper.
                    # NOTE: This returns tuple. First element has data.
                    nodeInfo = ZK_CLIENT.get("/group-membership/"+child)
                    nodeInfoJson = json.loads(nodeInfo[0])

                    handle_join(nodeInfoJson)            
        # Handle node failure.
        else:
            for member in CURRENT_MEMBERS:
                if member not in children:
                    logger.info ("Node: " + member + " has failed!")
                    CURRENT_MEMBERS.remove(member)
                    handle_failure(member)
                        
def handle_join(nodeInfo):
    # Create object to store in database.
    nodeToAdd = dict()
    nodeToAdd["name"] = nodeInfo["name"]
    nodeToAdd["nodeTemplate"] = nodeInfo["nodeTemplate"]
    nodeToAdd["status"] = "ACTIVE"

    interfaceToAdd = dict()
    interfaceToAdd["name"] = nodeInfo["interface"]
    interfaceToAdd["address"] = nodeInfo["address"]
    interfaceToAdd["network"] = nodeInfo["network"]

    nodeToAdd["interfaces"] = list()
    nodeToAdd["interfaces"].append(interfaceToAdd)

    nodeToAdd["processes"] = list()

    # Add node to database.
    mongoClient = MongoClient(MONGO_SERVER)
    db = mongoClient["ConfigSpace"]
    nColl = db["Nodes"]
    # NOTE: Update used with upsert instead of insert because
    # we might be adding node that had previously failed or
    # been removed.
    nColl.update({"name":nodeInfo["name"]}, nodeToAdd, upsert = True)

    # Check if any application already exists. If there are
    # applications then it means that the node join has to
    # be treated as hardware update and therefore the solver
    # should be invoked. If there are no applications then
    # this node is addition is happening at system initialization
    # time so do not invoke the solver.
    systemInitialization = True
    
    if "GoalDescriptions" in db.collection_names():
        gdColl = db["GoalDescriptions"]
        if gdColl.count() != 0:
            systemInitialization = False
    
    if not systemInitialization:
        logger.info ("New node added after system initialization. Inovking ManagementEngine to handle update.")
        invoke_management_engine(MONGO_SERVER, MANAGEMENT_ENGINE, True)
    else:
        logger.info ("New node added during system initialization.")

def handle_failure(node):
    logger.info ("Handeling failure of node: " + node)
    
    mongoClient = MongoClient(MONGO_SERVER)
    db = mongoClient["ConfigSpace"]
    nColl = db["Nodes"]
    ciColl = db["ComponentInstances"]

    # Mark node faulty in Nodes collection.
    nColl.update({"name": node, "status": "ACTIVE"}, 
                 {"$set": {"status": "FAULTY"}}, 
                 upsert=False)

    # Store names of affected component instances.
    findResults = nColl.find({"name": node, "status": "FAULTY"})
    failedComponentInstances = list()
    
    for findResult in findResults:
        failedNode = Serialize(**findResult)
        for p in failedNode.processes:
            process = Serialize(**p)
            for c in process.components:
                component = Serialize(**c)
                failedComponentInstances.append(component.name)

    # Update ComponentInstances collection using above collected information.
    for compInst in failedComponentInstances:
        ciColl.update({"name": compInst}, 
                      {"$set": {"status": "FAULTY"}})

    # Pull related processes in Nodes collection.
    nColl.update({"name": node, "status": "FAULTY"}, 
                 {"$pull": {"processes": {"name": {"$ne": "null"}}}})

    # Invoke solver for reconfiguration.
    invoke_management_engine(MONGO_SERVER, MANAGEMENT_ENGINE, False)

def main():
    global CURRENT_MEMBERS
    global ZK_CLIENT            # Zookeeper client object.
    global MANAGEMENT_ENGINE    # Management engine address.
    global MONGO_SERVER         # Mongo server address.
    
    monitoringServer = None
    
    argParser = argparse.ArgumentParser()

    argParser.add_argument("-c",
                           "--configFile",
                           help="CHARIOT configuration file to use.")

    args = argParser.parse_args()

    if args.configFile is None:
        logger.info ("Configuration file not provided.")
        argParser.print_help()
        argParser.exit()
    else:
        if (os.path.isfile(args.configFile)):
            logger.info ("Using configuration file: " + args.configFile)
            configParser = ConfigParser.ConfigParser()
            configParser.read(args.configFile)
            
            try:
                monitoringServer = configParser.get("Services", "MonitoringServer")
                logger.info ("Using monitoring server: " + monitoringServer)
            except ConfigParser.NoOptionError, ConfigParser.NoSectionError:
                monitoringServer = "localhost"
                logger.info ("Monitoring server cannot be extracted from configuration file. Using: " + monitoringServer)
            
            try:
                MONGO_SERVER = configParser.get("Services", "MongoServer")
                logger.info ("Using mongo server: " + MONGO_SERVER)
            except ConfigParser.NoOptionError, ConfigParser.NoSectionError:
                MONGO_SERVER = "localhost"
                logger.info ("Mongo server cannot be extracted from configuration file. Using: " + MONGO_SERVER)
    
            try:
                MANAGEMENT_ENGINE = configParser.get("Services", "ManagementEngine")
                logger.info ("Using management engine: " + MANAGEMENT_ENGINE)
            except ConfigParser.NoOptionError, ConfigParser.NoSectionError:
                MANAGEMENT_ENGINE = "localhost"
                logger.info ("Management Engine cannot be extracted from configuration file. Using: " + MANAGEMENT_ENGINE)
    
    CURRENT_MEMBERS = list()
    
    ZK_CLIENT = KazooClient(hosts=(monitoringServer+":2181"))
    
    # Add connection state listener to know the state
    # of connection between this client and ZooKeeper
    # server. 
    ZK_CLIENT.add_listener(connection_state_listener)
        
    # Start ZooKeeper client/server connection.
    ZK_CLIENT.start()
    
    # Create root group membership znode if it doesn't
    # already exist. 
    ZK_CLIENT.ensure_path("/group-membership")
    
    # Use watchers recipe to watch for changes in
    # children of group membership znode. Each 
    # child of this node is an ephemeral node that
    # represents a group member. We set send_event
    # to true and set membership_watch as the
    # corresponding callback function.
    ChildrenWatch(client = ZK_CLIENT,
                  path = "/group-membership",
                  func = membership_watch,
                  send_event = True)
    
    # Endless loop to ensure this detector process
    # doesn't die.    
    while True:
        time.sleep(5)

if __name__=='__main__':
    global logger
    logger = get_logger("chariot-nmw")
    main()
