#!/usr/bin/env python

import eventlet
eventlet.monkey_patch()

import datetime
import sys

from nova import rpc
from nova.rpc import amqp
from nova import log
from nova import flags
from nova import utils
from nova import exception
from nova.network.manager import NetworkManager

from ncclient import manager, operations


DELAY = 10
FLAGS = flags.FLAGS
LOG = log
RULESET_XML = """\
<config xmlns:xc="urn:ietf:params:xml:ns:netconf:base:1.0">
   <configuration><security><nat><static><rule-set>
        <name>{name}</name>
        {rules}
    </rule-set></static></nat></security></configuration>
</config>"""
RULE_XML = """\
<rule>
    <name>{name}</name>
    <static-nat-rule-match>
        <destination-address>
            <dst-addr>{destination}/32</dst-addr>
        </destination-address>
    </static-nat-rule-match>
    <then>
        <static-nat>
            <prefix>
                <addr-prefix>{prefix}/32</addr-prefix>
                <routing-instance>cloud</routing-instance>
            </prefix>
        </static-nat>
    </then>
</rule>"""
DELETE_XML = """\
<rule operation="delete">
    <name>{name}</name>
</rule>"""


def ensure_floating_forwards(ips):
    rules_xml = [
        RULE_XML.format(
            name=floating_ip.replace('.', '_'),
            destination=floating_ip,
            prefix=fixed_ip,
        )
        for context, floating_ip, fixed_ip, host, instance_id in ips
    ]
    apply_ruleset(rules_xml, default_operation="merge")


def remove_floating_forwards(ips):
    rules_xml = [
        DELETE_XML.format(name=address.replace('.', '_'))
        for context, address, instance_id in ips
    ]
    apply_ruleset(rules_xml, default_operation="none")


def apply_ruleset(rules_xml, default_operation):
    netconf_retries = FLAGS.netconf_retries
    netconf_ip = FLAGS.netconf_server_ip
    netconf_port = FLAGS.netconf_server_port
    netconf_user = FLAGS.netconf_server_user
    netconf_password = FLAGS.netconf_server_password
    netconf_ruleset = FLAGS.netconf_rule_set
    LOG.debug('Connecting to %s:%s', netconf_ip, netconf_port)
    for retry in xrange(netconf_retries):
        try:
            m = manager.connect_ssh(
                netconf_ip,
                port=netconf_port,
                unknown_host_cb=lambda host, fingerprint: True,
                username=netconf_user,
                password=netconf_password,
                allow_agent=False,
                look_for_keys=False,
            )
        except Exception as e:
            LOG.warning("Failed to connect %s:%s (%s)", netconf_ip, netconf_port, e)
        else:
            break
        eventlet.sleep(2**retry + 1)
    else:
        raise e
    with m as m:
        assert ':url' in m.server_capabilities
        LOG.debug('Starting private edit')
        for retry in xrange(netconf_retries):
            try:
                m.private()
            except operations.rpc.RPCError as e:
                if (e._message ==
                    '\nuncommitted changes will be discarded on exit\n'):
                    break
            LOG.debug('Failed to start private edit, try %d', retry)
            m.privateend()
            eventlet.sleep(2**retry + 1)
        else:
            LOG.debug('Failed to start private edit %r' % e)
            raise Exception("Failed to start private edit")
        ruleset_xml = RULESET_XML.format(
            name=netconf_ruleset,
            rules=''.join(rules_xml),
        )
        LOG.debug('Editing config')
        m.edit_config(
            target="candidate",
            config=ruleset_xml,
            default_operation=default_operation,
        )
        LOG.debug('Commiting')
        for retry in xrange(netconf_retries):
            try:
                reply = m.commit(confirmed=False)
            except operations.rpc.RPCError as e:
                pass
            else:
                for child in reply._root.getchildren():
                    error = child.find(
                        '{urn:ietf:params:xml:ns:netconf:base:1.0}rpc-error',
                    )
                    if error:
                        text = error.find(
                            '{urn:ietf:params:xml:ns:netconf:base:1.0}error-message',
                        )
                        LOG.warning('Commit error (%s)', text)
                        break
                else:
                    LOG.debug('Commited')
                    break
            LOG.warning('Failed to commit, try %d', retry)
            eventlet.sleep(2**retry + 1)
        else:
            LOG.error('Failed configuring network device %r.', netconf_ip)
            raise Exception("Failed configuring network device")


def associate_floating_ips(ips):
    if not ips:
        return
    ips = [
        (context, floating_ip, fixed_ip, host, instance_id)
        for fixed_ip, (context, floating_ip, host, instance_id) in {
            fixed_ip: (context, floating_ip, host, instance_id)
            for context, floating_ip, fixed_ip, host, instance_id in ips
        }.iteritems()
    ]
    try:
        # Change the device configuration
        LOG.info("Changing device configuration for %r", ips)
        ensure_floating_forwards(ips)
    except Exception as e:
        LOG.exception(e)
        return
    nm = NetworkManager()
    for context, floating_ip, fixed_ip, host, instance_id in ips:
        # Write the changes to the database
        LOG.info("Updating database for %s", fixed_ip)
        try:
            instance_ids = {instance_id}
            floating_ip_ref = nm.db.floating_ip_get_by_address(
                context,
                floating_ip,
            )
            fixed_ip_id = floating_ip_ref['fixed_ip_id']
            if fixed_ip_id:
                old_fixed_ip_ref = nm.db.fixed_ip_get(context, fixed_ip_id)
                if old_fixed_ip_ref and old_fixed_ip_ref['instance_id']:
                    instance_ids.add(old_fixed_ip_ref['instance_id'])
            nm.db.floating_ip_fixed_ip_associate(
                context,
                floating_ip,
                fixed_ip,
                host,
            )
            for id_ in instance_ids:
                instance = nm.db.instance_get(context, id_)
                nm.get_instance_nw_info(
                    context.elevated(),
                    instance['id'],
                    instance['uuid'],
                    instance['instance_type_id'],
                    instance['host'],
                )
        except Exception as e:
            LOG.exception(e)


def disassociate_floating_ips(ips):
    if not ips:
        return
    ips = [
        (context, address, instance_id)
        for address, (context, instance_id) in {
            address: (context, instance_id)
            for context, address, instance_id in ips
        }.iteritems()
    ]
    nm = NetworkManager()
    good_ips = []
    for context, address, instance_id in ips:
        floating_ip_ref = nm.db.floating_ip_get_by_address(context, address)
        if floating_ip_ref['fixed_ip_id']:
            good_ips.append((context, address, instance_id))
        else:
            LOG.warning("Address %s already removed!" % address)
    if not good_ips:
        return
    try:
        # Change the device configuration
        LOG.info("Changing device configuration for %r", ips)
        remove_floating_forwards(good_ips)
    except Exception as e:
        LOG.exception(e)
        return
    for context, address, instance_id in good_ips:
        # Write the changes to the database
        LOG.info("Updating database for %s", address)
        try:
            nm.db.floating_ip_disassociate(context, address)
            try:
                instance = nm.db.instance_get(context, instance_id)
                nm.get_instance_nw_info(
                    context.elevated(),
                    instance['id'],
                    instance['uuid'],
                    instance['instance_type_id'],
                    instance['host'],
                )
            except exception.InstanceNotFound:
                pass
        except Exception as e:
            LOG.exception(e)


def on_message(data):
    method = data.get('method')
    kwargs = data.get('args', {})
    context = amqp.unpack_context(data)
    if method == 'srx_associate_floating_ip':
        assoc_acc.add((
            context,
            kwargs.get('floating_ip'),
            kwargs.get('fixed_ip'),
            kwargs.get('host'),
            kwargs.get('instance_id'),
        ))
    elif method == 'srx_disassociate_floating_ip':
        dis_acc.add((
            context,
            kwargs.get('address'),
            kwargs.get('instance_id'),
        ))


class Accumulator(object):
    """Accumulates entries until the oldest one is older than timeout."""

    def __init__(self, timeout):
        self.cache = []
        self.oldest = None
        self.timeout = datetime.timedelta(seconds=timeout)

    def add(self, item):
        if not self.cache:
            self.oldest = datetime.datetime.now()
        self.cache.append(item)

    def purge(self):
        content = self.cache
        self.cache = []
        self.oldest = None
        return content

    def check(self):
        if self.oldest and self.oldest + self.timeout < datetime.datetime.now():
            return self.purge()
        return []

    def delay(self):
        if not self.oldest:
            return 0
        return (datetime.datetime.now() - self.oldest).seconds



if __name__ == '__main__':
    utils.default_flagfile()
    FLAGS(sys.argv)
    log.setup()
    assoc_acc = Accumulator(DELAY)
    dis_acc = Accumulator(DELAY)
    for retry in xrange(FLAGS.netconf_retries):
        try:
            connection = rpc.create_connection()
        except Exception as e:
            log.warning("Can't connect, retry %d", retry + 1)
            eventlet.sleep(2**retry + 1)
        else:
            break
    else:
        log.error("Failed to connect.")
        raise e
    connection.declare_topic_consumer('srxproxy', on_message)
    connection.declare_topic_consumer('srxproxy.%s' % flags.FLAGS.host, on_message)
    connection.consume_in_thread()
    while True:
        associate_floating_ips(assoc_acc.check())
        disassociate_floating_ips(dis_acc.check())
        delay = min(DELAY - assoc_acc.delay(), DELAY - dis_acc.delay())
        eventlet.sleep(delay)