import sys
import json, ast
from RTEInterface import RTEInterface
from kafka import KafkaConsumer

consumer = KafkaConsumer('deployed', 'removeWorker',
                         "request", bootstrap_servers='10.129.6.5:9092')

RTEInterface.Connect('thrift', "10.129.2.201", 20206)
tableId = "ingress::dispatch"
ruleDictionary = {}

def makeRule(ip, port, mac, functionHash, tableId, rule_name, default_rule):
    actions = '{  "type" : "ingress::dispatch_act",  "data" : {  "dstAddr" : { "value" : "%s" }, \
        "dstPort" : { "value" : "%d" } , "egress_port": { "value": "v0.1" }, "ethernetAddr": { "value": "%s" } } }' \
         % (ip, int(port), mac)
    match = '{ "map_hdr.function_id" : {  "value" : %d} } ' % (functionHash)
    rule = {
        "tableId": tableId,
        "rule_name": rule_name,
        "default_rule": default_rule,
        "match": match,
        "actions": actions  
    }
    return rule

def addRule(worker):
    functionHash = worker[u'functionHash']
    rule_name = "dispatch_to_worker" + functionHash
    functionHash = int(functionHash[0:5], 16)
    ip = str(worker[u'node_id']).strip()
    port = int(worker[u'portExternal'])
    mac = str(worker[u'mac']).strip()
    default_rule = False
    rule = makeRule(ip, port, mac, functionHash, tableId, rule_name, default_rule)
    ruleDictionary[functionHash] = rule
    print "rule added: ", ruleDictionary
    RTEInterface.Tables.AddRule(
        rule["tableId"], rule["rule_name"], rule["default_rule"], rule["match"], rule["actions"])
    ruleList = RTEInterface.Tables.ListRules(tableId)
    print "new rule list: ", ruleList, "\n\n"
    return 0

def deleteRule(worker):
    functionHash = worker[u'functionHash']
    functionHash = int(functionHash[0:5], 16)
    rule = ruleDictionary[functionHash]
    RTEInterface.Tables.DeleteRule(
        rule["tableId"], rule["rule_name"], rule["default_rule"], rule["match"], rule["actions"])
    del ruleDictionary[functionHash]
    ruleList = RTEInterface.Tables.ListRules(tableId)
    print "deleted rule"
    print "new rule list: ", ruleList, "\n\n"
    
    return 0

for msg in consumer:
    if msg.topic == "deployed":
        msg = msg.value.decode('utf-8')
        worker = json.loads(msg)
        addRule(worker)
    
    elif msg.topic == "removeWorker":
        msg = msg.value.decode('utf-8')
        worker = json.loads(msg)
        deleteRule(worker)
