"""
Copyright 2017 Pani Networks Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

"""

#
# Functions dealing with VPC.
#

import datetime
import logging
import random

import boto.vpc
import boto.utils

from vpcrouter.errors       import VpcRouteSetError
from vpcrouter.currentstate import CURRENT_STATE
from vpcrouter.utils        import is_cidr_in_cidr


def get_ec2_meta_data():
    """
    Get meta data about ourselves, if we are on an EC2 instance.

    In particular, this returns the VPC ID and region of this instance.

    If we are not on an EC2 instance it returns an empty dict.

    """
    # The timeout is just for the connection attempt, but between retries there
    # is an exponential back off in seconds. So, too many retries and it can
    # possibly block for a very long time here. Main contributor to waiting
    # time here is the number of retries, rather than the timeout time.
    try:
        md     = boto.utils.get_instance_metadata(timeout=2, num_retries=2)
        vpc_id = md['network']['interfaces']['macs'].values()[0]['vpc-id']
        region = md['placement']['availability-zone'][:-1]
        return {"vpc_id" : vpc_id, "region_name" : region}
    except:
        # Any problem while getting the meta data? Assume we are not on an EC2
        # instance.
        return {}


def connect_to_region(region_name):
    """
    Establish connection to AWS API.

    """
    logging.debug("Connecting to AWS region '%s'" % region_name)
    con = boto.vpc.connect_to_region(region_name)
    if not con:
        raise VpcRouteSetError("Could not establish connection to "
                               "region '%s'." % region_name)
    return con


def _make_ip_subnet_lookup(vpc_info):
    """
    Updates the vpc-info object with a lookup for IP -> subnet.

    """
    # We create a reverse lookup from the instances private IP addresses to the
    # subnets they are associated with. This is used later on in order to
    # determine whether routes should be set in an RT: Is the RT's subnet
    # associated with ANY of the IP addresses in the route spec? To make this
    # easy, we collect the IPs and subnets of all EC2 instances in the VPC.
    # Once we get a route spec, we create a list of subnets for only the
    # cluster nodes. The assumption is that not all EC2 instances in the VPC
    # necessarily belong to the cluster. We really want to narrow it down to
    # the cluster nodes only. See make_cluster_node_subnet_list().
    vpc_info['ip_subnet_lookup'] = {}
    for instance in vpc_info['instances']:
        for interface in instance.interfaces:
            subnet_id = interface.subnet_id
            for priv_addr in interface.private_ip_addresses:
                vpc_info['ip_subnet_lookup'][priv_addr.private_ip_address] = \
                                                                    subnet_id


def get_vpc_overview(con, vpc_id, region_name):
    """
    Retrieve information for the specified VPC.

    If no VPC ID was specified then just pick the first VPC we find.

    Returns a dict with the VPC's zones, subnets, route tables and
    instances.

    """
    logging.debug("Retrieving information for VPC '%s'" % vpc_id)
    d = {}
    d['zones'] = con.get_all_zones()

    # Find the specified VPC, or just use the first one
    all_vpcs    = con.get_all_vpcs()
    if not all_vpcs:
        raise VpcRouteSetError("Cannot find any VPCs.")

    if not vpc_id:
        # Just grab the first available VPC and use it, if no VPC specified
        vpc = all_vpcs[0]
        vpc_id = vpc.id
    else:
        # Search through the list of VPCs for the one with the specified ID
        vpc = None
        for v in all_vpcs:
            if v.id == vpc_id:
                vpc = v
                break
        if not vpc:
            raise VpcRouteSetError("Cannot find specified VPC '%s' "
                                   "in region '%s'." % (vpc_id, region_name))
    d['vpc'] = vpc

    vpc_filter = {"vpc-id" : vpc_id}  # Will use this filter expression a lot

    # Now find the subnets, route tables and instances within this VPC
    d['subnets']      = con.get_all_subnets(filters=vpc_filter)
    d['route_tables'] = con.get_all_route_tables(filters=vpc_filter)

    # Route tables are associated with subnets. Maintain a lookup table from
    # RT ID to (list of) subnets. This is necessary later on, because
    # we only want to set a route in an RT if at least one of the ENIs we are
    # dealing with is associated with the RT. That way, we don't have to set
    # the route in all tables all the time.
    d['rt_subnet_lookup'] = {}
    for rt in d['route_tables']:
        for assoc in rt.associations:
            if hasattr(assoc, 'subnet_id'):
                subnet_id = assoc.subnet_id
                if subnet_id:
                    d['rt_subnet_lookup'].setdefault(rt.id, []). \
                                                            append(subnet_id)

    # Get all the cluster instances (all EC2 instances associated with this
    # VPC).
    reservations   = con.get_all_reservations(filters=vpc_filter)
    d['instances'] = []
    for r in reservations:  # a reservation may have multiple instances
        d['instances'].extend(r.instances)

    # Add reverse lookup from the instances private IP addresses to the
    # subnets they are associated with. More info in the doc for that function.
    _make_ip_subnet_lookup(d)

    # Maintain a quick instance lookup for convenience
    d['instance_by_id'] = {}
    for i in d['instances']:
        d['instance_by_id'][i.id] = i

    return d


def find_instance_and_eni_by_ip(vpc_info, ip):
    """
    Given a specific IP address, find the EC2 instance and ENI.

    We need this information for setting the route.

    Returns instance and emi in a tuple.

    """
    for instance in vpc_info['instances']:
        for eni in instance.interfaces:
            for pa in eni.private_ip_addresses:
                if pa.private_ip_address == ip:
                    return instance, eni
    raise VpcRouteSetError("Could not find instance/eni for '%s' "
                           "in VPC '%s'." % (ip, vpc_info['vpc'].id))


def get_instance_private_ip_from_route(instance, route):
    """
    Find the private IP and ENI of an instance that's pointed to in a route.

    Returns (ipaddr, eni) tuple.

    """
    ipaddr = None
    for eni in instance.interfaces:
        if eni.id == route.interface_id:
            ipaddr = eni.private_ip_address
            break
    return ipaddr, eni if ipaddr else None


def _choose_different_host(old_ip, ip_list, failed_ips, questionable_ips):
    """
    Randomly choose a different host from a list of hosts.

    Pick from fully healthy IPs first (neither failed nor questionable).
    If we don't have any of those, pick from questionable ones next.

    If no suitable hosts can be found in the list (if it's empty or all hosts
    are in the failed_ips list) it will return None.

    The old IP (if any) is passed in. We will try to avoid returning this same
    old IP under the right circumstances. If no old IP is known, None can be
    passed in for it instead.

    """
    if not ip_list:
        # We don't have any hosts to choose from.
        return None

    ip_set           = set(ip_list)
    failed_set       = set(failed_ips)
    # Consider only those questionable IPs that aren't also failed and make
    # sure all of the ones in the questionable list are at least also present
    # in the overall IP list.
    questionable_set = set(questionable_ips).intersection(ip_set). \
                                             difference(failed_set)

    # Get all healthy IPs that are neither failed, nor questionable
    healthy_ips = list(ip_set.difference(failed_set, questionable_set))

    if healthy_ips:
        # Return one of the completely healthy IPs
        return random.choice(healthy_ips)

    if questionable_set:
        # Don't have any completely healthy ones, so return one of the
        # questionable ones. Not perfect, but at least may still provide
        # routing functionality for some time.
        if old_ip not in questionable_set:
            # We may be here because the original address was questionable. If
            # only other questionable ones are available then there's no point
            # changing the address. We only change if the old address wasn't
            # one of the questionable ones already.
            return random.choice(list(questionable_set))

    # We got nothing...
    return None


def _rt_state_update(route_table_id, dcidr, router_ip="(none)",
                     instance_id="(none)", eni_id="(none)",
                     old_router_ip="(none)", msg="(none)"):
    """
    Store a message about a VPC route in the current state.

    """
    buf = "inst: %s, eni: %s, r_ip: %-15s, o_r_ip: %-15s, msg: %s" % \
          (instance_id, eni_id, router_ip, old_router_ip, msg)
    CURRENT_STATE.vpc_state.setdefault('route_tables', {}). \
                            setdefault(route_table_id, {})[dcidr] = buf


def _update_route(dcidr, router_ip, old_router_ip,
                  vpc_info, con, route_table_id, update_reason):
    """
    Update an existing route entry in the route table.

    """
    instance = eni = None
    try:
        instance, eni = find_instance_and_eni_by_ip(vpc_info, router_ip)

        logging.info("--- updating existing route in RT '%s' "
                     "%s -> %s (%s, %s) (old IP: %s, reason: %s)" %
                     (route_table_id, dcidr, router_ip,
                      instance.id, eni.id, old_router_ip, update_reason))
        try:
            con.replace_route(