'''
Discovery Service main class
Created on Oct 19, 2016

@author: riaps
'''
import time
import sys
import re
import os

import zmq
from czmq import Zsys
import capnp

from .dbase_redis import RedisDbase
from .dbase_dht import DhtDbase
from riaps.utils.config import Config
from riaps.proto import disco_capnp
from riaps.consts.defs import *
from riaps.utils.ifaces import getNetworkInterfaces
from riaps.run.exc import DatabaseError
import logging

class DiscoService(object):
    '''
    Discovery service main class. 
    '''
    def __init__(self,dbaseLoc):
        self.logger = logging.getLogger(__name__)
        self.setupIfaces()
        self.suffix = self.macAddress
        czmq_ctx = Zsys.init()
        self.context = zmq.Context.shadow(czmq_ctx.value)
        Zsys.handler_reset() 
        self.context = zmq.Context()
        self.dbase = None 
        if Config.DISCO_TYPE == 'redis':
            self.dbase = RedisDbase(self.context,dbaseLoc)
        elif Config.DISCO_TYPE == 'opendht':
            self.dbase = DhtDbase(self.context,self.hostAddress,dbaseLoc)
        else:
            pass    
        self.registrations = { }            # Registered 'producers': app.actor -> { producer* }
        self.clients = { }                  # Information about client actors
        self.enrollments = { }              # Registered 'consumers': app.actor -> { consumer* } 
        self.devices = set()                # Local devices
        self.GROUPMSG_RE = re.compile("\w*\@\w*\.\w*")
    
    def setupIfaces(self):
        '''
        Find the IP addresses of the (host-)local and network(-global) interfaces
        '''
        (globalIPs,globalMACs,globalNames,_localIP) = getNetworkInterfaces()
        try:
            assert len(globalIPs) > 0 and len(globalMACs) > 0
        except:
            self.logger.error("Error: no active network interface")
            raise
        globalIP = globalIPs[0]
        globalMAC = globalMACs[0]
        if Config.NIC_NAME != globalNames[0]:
            Config.NIC_NAME = globalNames[0]
        self.hostAddress = globalIP
        self.macAddress = globalMAC
        
    def start(self):
        self.logger.info("starting")
        self.server = self.context.socket(zmq.REP)              # Create main server socket for client requests
        endpoint = const.discoEndpoint
        self.server.bind(endpoint)
        
        self.dbase.start()                                      # Start database
        time.sleep(0.0001)                                      # Yield to database so that it can start
        
        self.poller = zmq.Poller()                              # Set up initial poller (only on the main server socket)  
        self.poller.register(self.server,zmq.POLLIN)

        self.clientUpdates = []
    
    def run(self):
        '''
        Main loop of the discovery service
        '''
        self.logger.info("running")
        while 1:
            self.clientUpdates = []
            sockets = dict(self.poller.poll(1000.0))            # Poll client messages, with timeout 1 sec
            if len(sockets) == 0:                               # If no message but timeout expired, 
                try: 
                    self.clientUpdates = self.dbase.fetchUpdates()  # then fetch updates from database
                except DatabaseError:
                    self.logger.info("reconnecting database")
                    self.dbase.start()
                    time.sleep(0.0001)     
            elif self.server in sockets:                        # else check if there is a server request, handle it 
                msg = self.server.recv()
                self.handle(msg)
                del sockets[self.server]
            else:
                pass
            for note in self.clientUpdates:                     # Handle all client updates (received from the 
                self.handleUpdate(note)                         # database and generated by the server requests

    def appActorName(self,appName,appActorName):
        return "%s/%s" % (appName,appActorName) 
    
    def isDeviceName(self,appName,actorName):
        return ("/" + appName + '/' + actorName + "/") in self.devices
        
    def setupClient(self,appName,appVersion,actorName,isDevice):
        '''
        Set up a new client of the discovery service. The client actors are to register with
        the service using the 'server' (REQ/REP) socket. The service will then create a dedicated
        (PAIR) socket for the client to connect to. This socket is used as a private communication
        channel between a specific client actor and the service.   
        '''
        sock = self.context.socket(zmq.PAIR)
        port = sock.bind_to_random_port('tcp://127.0.0.1')
        clientKeyBase = "/" + appName + '/' + actorName + "/"
        self.clients[clientKeyBase] = sock
        if isDevice: self.devices.add(clientKeyBase)
        clientKeyLocal = clientKeyBase + self.macAddress
        self.clients[clientKeyLocal] = port
        clientKeyGlobal = clientKeyBase + self.hostAddress
        self.clients[clientKeyGlobal] = port
        actorKey = self.appActorName(appName,actorName)
        self.registrations[actorKey] = set()
        self.clients[actorKey] = []
        self.enrollments[actorKey] = set()
        return port
    
    def unsetupClient(self,appName,appVersion,actorName):
        '''
        Remove a client   
        '''
        clientKeyBase = "/" + appName + '/' + actorName + "/"
        if clientKeyBase not in self.clients:
            return -1                                   # Client has already been removed
        sock = self.clients[clientKeyBase]
        clientKeyLocal = clientKeyBase + self.macAddress
        assert clientKeyLocal in self.clients
        port = self.clients[clientKeyLocal]
        clientKeyGlobal = clientKeyBase + self.hostAddress
        assert clientKeyGlobal in self.clients
        # port = self.clients[clientKeyGlobal]         # Must have the same value
        self.logger.info("unsetupClient - unbind/close sock")
        sock.unbind('tcp://127.0.0.1:' + str(port))
        sock.close(0)
        client = ""
        # Remove all services registered by this client  
        del self.clients[clientKeyBase]
        del self.clients[clientKeyLocal]
        del self.clients[clientKeyGlobal]
        if clientKeyBase in self.devices:
            self.devices.discard(clientKeyBase)
        actorKey = self.appActorName(appName,actorName)
        # assert actorKey in self.registrations
        for (key,value) in self.registrations.get(actorKey,set()):
            self.logger.info("unsetupClient - dbase remove %r.%r.%r" % (actorName,key,value))
            self.dbase.remove(key, value)               # Remove k/v pair from database
        # assert actorKey in self.enrollments
        for (key,target) in self.enrollments.get(actorKey,set()):
            self.logger.info("unsetupClient - dbase detach %r.%r.%r" % (actorName,key,target))
            self.dbase.detach(key, target)              # Detach target from key
        del self.registrations[actorKey]
        del self.enrollments[actorKey]
        return port
    
    def handleActorReg(self,msg):
        '''
        Handle the registration of an application actor with the service. 
        '''
        actReg = msg.actorReg
        appName = actReg.appName
        appVersion = actReg.version   
        actorName = actReg.actorName
        isDevice = actReg.isDevice
         
        self.logger.info("handleActorReg: %s,%s,%s,%s" % (appName, appVersion, actorName, "[device]" if isDevice else ""))
        
        clientPort = self.setupClient(appName,appVersion,actorName,isDevice)
        
        # Option: store in db host.app.vers.actr -> port? 
        
        rsp = disco_capnp.DiscoRep.new_message()
        rspMessage = rsp.init('actorReg')
        rspMessage.status = "ok"
        rspMessage.port = clientPort
        rspBytes = rsp.to_bytes()
        self.server.send(rspBytes)
    
    def handleActorUnreg(self,msg):
        '''
        Handle the unregistration of an application actor with the service. 
        '''
        actReg = msg.actorUnreg
        appName = actReg.appName
        appVersion = actReg.version   
        actorName = actReg.actorName
         
        self.logger.info("handleActorUnreg: %s,%s,%s" % (appName, appVersion, actorName))
        
        clientPort = self.unsetupClient(appName,appVersion,actorName)
        
        # Option: store in db host.app.vers.actr -> port? 
        
        rsp = disco_capnp.DiscoRep.new_message()
        rspMessage = rsp.init('actorUnreg')
        rspMessage.status = "ok"
        rspMessage.port = clientPort
        rspBytes = rsp.to_bytes()
        self.server.send(rspBytes)
        
        self.logger.info("handleActorUnreg: %s,%s,%s - completed" % (appName, appVersion, actorName))
        
    def buildInsertKeyValuePair(self,appName,appActorName,msgType,kind,scope,host,port,):
        '''
        Construct a database key,value pair to be used when a service is registered.
        The key allows the identifying the specific object in the actor the service
        is provided through. The value is a host:port pair.
        '''
        key = '/' + appName + '/' + msgType + '/' + kind
        if scope == "global":
            pass
        elif scope == "local":
            key = key + ":" + str(self.macAddress)
        elif scope == "internal":
            if self.isDeviceName(appName, appActorName):
                actorName,_rest = appActorName.split('.')
            else:
                actorName = appActorName
            key = key + ":" + str(self.macAddress) + ":" + actorName
        value = str(host) + ':' + str(port)
        return (key,value)
            
    def buildLookupKey(self,appName,msgType,kind,scope,clientHost,clientActorName,clientInstanceName,clientPortName):
        '''
        Construct a key used to lookup a service. Construct also a string that identifies the client of the lookup 
        '''
        # Map the requestor's kind into the provider's kind
        kindMap = { "sub" : "pub" , "clt" : "srv", "req" : "rep", "rep" : "req", "qry" : "ans", "gsub" : "gpub"} 
        key = '/' + appName + '/' + msgType + '/' + kindMap[kind]
        if scope == "global":
            pass
        elif scope == "local":                      # If the request is host-local, add the mac address to the end. Used to 
            key = key + ":" + str(self.macAddress)  # distinguish node-specific local requests. (The database is shared!)
        elif scope == "internal":                    # If the request is actor-internal, add the actor name at the end. 
            if self.isDeviceName(appName, clientActorName):
                actorName,_rest = clientActorName.split('.')
            else:
                actorName = clientActorName
            key = key + ":" + str(self.macAddress) + ":" + actorName 
        client = '/' + appName + '/' + clientActorName + '/' + clientHost + '/' + clientInstanceName + '/' + clientPortName 
        if scope == "global":
            pass
        elif scope == "local":
            client = client + ":" + str(self.macAddress)
        elif scope == "internal":
            client = client + ":" + str(self.macAddress) + ":" + clientActorName 
        return (key,client)

    def handleServiceReg(self,msg):
        '''
        Handle the service registration message
        '''
        reqMsg = msg.serviceReg                             # Parse the message
        path = reqMsg.path
        appName = path.appName 
        actorName = path.actorName
        msgType = path.msgType
        kind = str(path.kind)
        scope = str(path.scope)
        
        socket = reqMsg.socket
        host = socket.host
        port = socket.port
        
        self.logger.info("handleServiceReg: %s,%s,%s,%s,%s,%s,%s" % (appName,actorName,msgType,kind,scope,host,port))

        regKey = self.appActorName(appName, actorName)
        
        # if regKey not in self.registrations:
        (key,value) = self.buildInsertKeyValuePair(appName, actorName, msgType, kind, scope,host, port)
        clients = self.dbase.insert(key,value)
        self.logger.info("handleServiceReg: clients = %r",clients)
        self.registrations[regKey] |= {(key,value)} 
        
        rep = disco_capnp.DiscoRep.new_message()            # Construct response
        repMsg = rep.init('serviceReg')
        repMsg.status = "ok"
        repBytes = rep.to_bytes()
        self.server.send(repBytes)
        
        if len(clients) != 0:
            self.clientUpdates += [(key,value,clients)]
    
    def handleServiceLookup(self,msg):
        '''
        Handle a service lookup message
        '''
        reqMsg = msg.serviceLookup                          # Parse the message
        path = reqMsg.path
        appName = path.appName
        msgType = path.msgType
        kind = str(path.kind)
        scope = str(path.scope)
        client = reqMsg.client
        clientActorHost = client.actorHost
        clientActorName = client.actorName
        clientInstanceName = client.instanceName
        clientPortName = client.portName
        client = (appName,clientActorHost,clientActorName,clientInstanceName,clientPortName)


        (key,client) = self.buildLookupKey(appName, msgType, kind, scope,
                                           clientActorHost, clientActorName, 
                                           clientInstanceName,clientPortName)
        result = self.dbase.fetch(key,client)
        
        regKey = self.appActorName(appName, clientActorName)
        self.enrollments[regKey] |=  {(key,client)}
        
        self.logger.info("handleServiceLookup: clients = %r",self.enrollments[regKey])
        self.logger.info("handleServiceLookup:%s,%s,%s,%s,%s,%s -> %r"
                           % (appName,str(client),msgType,kind,scope,clientInstanceName,result))
           
        rep = disco_capnp.DiscoRep.new_message()            # Construct the response: all providers of the requested service
        repMsg = rep.init('serviceLookup')
        repMsg.status = "ok"
        sockets = repMsg.init('sockets',len(result))
        i = 0
        for elt in result:
            pair = str(elt).split(':')
            sockets[i].host = str(pair[0])
            sockets[i].port = int(pair[1])
            i += 1
        repBytes = rep.to_bytes()
        self.server.send(repBytes)
    
    def handleServiceUnreg(self,msg):
        '''
        Handle the service un-registration message
        '''
        reqMsg = msg.serviceUnreg                             # Parse the message
        path = reqMsg.path
        appName = path.appName 
        actorName = path.actorName
        msgType = path.msgType
        kind = str(path.kind)
        scope = str(path.scope)
        
        socket = reqMsg.socket
        host = socket.host
        port = socket.port
        
        self.logger.info("handleServiceUnreg: %s,%s,%s,%s,%s,%s,%s" % (appName,actorName,msgType,kind,scope,host,port))

        regKey = self.appActorName(appName, actorName)

        (key,value) = self.buildInsertKeyValuePair(appName, actorName, msgType, kind, scope,host, port)
        self.dbase.remove(key,value)
        self.registrations[regKey] -= {(key,value)}
        
        rep = disco_capnp.DiscoRep.new_message()            # Construct response
        repMsg = rep.init('serviceUnreg')
        repMsg.status = "ok"
        repBytes = rep.to_bytes()
        self.server.send(repBytes)
        
        # TODO: there should be a 'negative' update for ports (disconnect) 
        #if len(clients) != 0:
        #    self.clientUpdates = [(key,(value),clients)]
    
#     def makeGroupMessageType(self,msgType,groupType,groupName):
#         return str(msgType) + '@' + str(groupType) + '.' + str(groupName)
#     
#     def isGroupMessageType(self,msg):
#         return self.GROUPMSG_RE.match(msg)
#     
#     def parseGroupMessageType(self,msgName):
#         msgType,rest = msgName.split('@')
#         groupType,groupName = rest.split('.')
#         return (msgType,groupType,groupName)    

    def handleServiceUnlookup(self,msg):
        '''
        Handle a service unlookup message
        '''
        reqMsg = msg.serviceUnlookup                          # Parse the message
        path = reqMsg.path
        appName = path.appName
        msgType = path.msgType
        kind = str(path.kind)
        scope = str(path.scope)
        client = reqMsg.client
        clientActorHost = client.actorHost
        clientActorName = client.actorName
        clientInstanceName = client.instanceName
        clientPortName = client.portName
        client = (appName,clientActorHost,clientActorName,clientInstanceName,clientPortName)


        (key,client) = self.buildLookupKey(appName, msgType, kind, scope,
                                           clientActorHost, clientActorName, 
                                           clientInstanceName,clientPortName)
        
        regKey = self.appActorName(appName, clientActorName)
        self.dbase.detach(key,client)
        self.enrollments[regKey] -=  {(key,client)}
        
        self.logger.info("handleServiceUnlookup:%s,%s,%s,%s,%s,%s"
                           % (appName,str(client),msgType,kind,scope,clientInstanceName))   
        rep = disco_capnp.DiscoRep.new_message()            # Construct the response: all providers of the requested service
        repMsg = rep.init('serviceUnlookup')
        repMsg.status = "ok"
        repBytes = rep.to_bytes()
        self.server.send(repBytes)

    def handleGroupJoin(self,msg):
        '''
        Handle the case when a component wants to join a group
        '''
        self.logger.info("handleGroupJoin: %s",str(msg))
        groupJoin = msg.groupJoin
        _appName = groupJoin.appName
        groupId = groupJoin.groupId
        _groupType = groupId.groupType
        _groupName = groupId.groupName 
        _services = groupJoin.services
        _componentId = groupJoin.componentId
        _pid = groupJoin.pid
        
#         for s in services:
#             messageType = s.messageType
#             address = s.address
#             (host,port) = address.split(':') 
#             
#             regKey = self.appActorName(appName, 'group')
#             clients = []
#         
#             # if regKey not in self.registrations:
#             groupMessageType = self.makeGroupMessageType(messageType,groupType,groupName)
#             (key,value) = self.buildInsertKeyValuePair(appName, groupMessageType, "gpub", "global",host, port)
#             clients = self.dbase.insert(key,value)
#             self.registrations[regKey].append((key,value))
#             
#         
        rsp = disco_capnp.DiscoRep.new_message()
        rspMessage = rsp.init('groupJoin')
        rspMessage.status = "ok"

        rspBytes = rsp.to_bytes()
        self.server.send(rspBytes)
        
    def handle(self,msgBytes):
        '''
        Dispatch the request based on the message type
        '''
        msg = disco_capnp.DiscoReq.from_bytes(msgBytes)
        which = msg.which()
        if which == 'actorReg':
            self.handleActorReg(msg)
        elif which == "serviceReg":
            self.handleServiceReg(msg)
        elif which == "serviceLookup":
            self.handleServiceLookup(msg)
        elif which == 'actorUnreg':
            self.handleActorUnreg(msg)
        elif which == 'serviceUnlookup':
            self.handleServiceUnlookup(msg)
        elif which == 'serviceUnreg':
            self.handleServiceUnreg(msg)
        elif which == 'groupJoin':
            self.handleGroupJoin(msg)
        else:
            pass
        
    def handleUpdate(self,msg):
        '''
        Handle a notification message received from the database.
        The notification triggers the notification of client actors about the new service provider
        '''
        self.logger.info("handleNote: %s",str(msg))
        (_key,value,clients) = msg                       # Parse notification message
        pair = re.split(':',value)
        host = pair[0]
        port = pair[1]
        for client in clients:                          # For each client:
            clientString = client                       # .decode('utf-8')           # Parse the client string
            spl = re.split('/',clientString)
            _skip = spl[0]
            appName = spl[1]
            actorName = spl[2]
            actorHost = spl[3]
            instanceName = spl[4]
            portMacAct = re.split(':',spl[5])
            portName = portMacAct[0]
            macAddr,actName = None, None
            lenPMA = len(portMacAct)
            if lenPMA == 1:
                pass
            elif lenPMA == 2:
                macAddr = portMacAct[1]
            elif lenPMA == 3:
                macAddr = portMacAct[1]
                actName = portMacAct[2]
            clientKeyBase = "/" + appName + '/' + actorName + "/"
            if clientKeyBase not in self.clients:
                continue
            if self.hostAddress != actorHost:
                continue
            clientSocket = self.clients[clientKeyBase]

            updMsg = disco_capnp.DiscoUpd.new_message() # Construct the notification message to be sent to the client actor
            updMsgPortUpd = updMsg.portUpdate
            updMsgClient = updMsgPortUpd.client
            updMsgClient.actorHost = actorHost
            updMsgClient.actorName = actorName
            updMsgClient.instanceName = instanceName
            updMsgClient.portName = portName
            updMsgPortUpd.scope = 'global' if macAddr == None else 'local' if actName == None else 'internal'
            updMsgPortUpd.socket.host = host 
            updMsgPortUpd.socket.port = int(port)
            
            msgBytes = updMsg.to_bytes()
            self.logger.info("send update to actor %s.%s.%s:%s %s:%s" 
                             % (actorHost,actorName,instanceName,portName,str(host),str(port)))
            clientSocket.send(msgBytes)                     # Send message to client 

    def terminate(self):
        self.logger.info("terminating")
        # Clean up everything
        self.dbase.terminate()
        self.context.destroy()
        time.sleep(1.0)
        self.logger.info("terminated")
        os._exit(0)

    