# Copyright 2013: Mirantis Inc.
# All Rights Reserved.
#
#    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.

import os
import re
import time

import netaddr
import six
from six import moves

from rally.common.i18n import _
from rally.common import log as logging
from rally.common import utils
from rally.deploy.serverprovider import provider
from rally import exceptions


LOG = logging.getLogger(__name__)
INET_ADDR_RE = re.compile(r' *inet ((\d+\.){3}\d+)\/\d+ .*')
IPT_PORT_TEMPLATE = ('iptables -t nat -{action} PREROUTING -d {host_ip}'
                     ' -p tcp --syn --dport {port}'
                     ' -j DNAT --to-destination {ip}:22')


def _get_script(filename):
    path = os.path.abspath(os.path.join(os.path.dirname(__file__),
                           'lxc', filename))
    return open(path, 'rb')


def _get_script_from_template(template_filename, **kwargs):
    template = _get_script(template_filename).read()
    return moves.StringIO(template.format(**kwargs))


class LxcHost(object):
    """Represent lxc enabled host."""

    def __init__(self, server, config):
        """Initialize LxcHost object.

        :param server:  Server object
        :param config:  dictionary with following key/values:
            network         ipv4 network for containers
            lxc_bridge      bridge interface name (default lxcbr0)
            tunnel_to       ip address for make tunnel to
            forward_ssh     use ssh port forwarding (do not use for
                            controller nodes)

        """
        self.config = config
        if 'network' in config:
            self.network = netaddr.IPNetwork(config['network'])
        else:
            self.network = None
        self.server = server
        self.containers = []
        self.path = '/var/lib/lxc/'
        self._port_cache = {}

    def _get_updated_server(self, **kwargs):
        credentials = self.server.get_credentials()
        credentials.update(kwargs)
        return provider.Server.from_credentials(credentials)

    @property
    def backingstore(self):
        if not hasattr(self, '_backingstore'):
            code = self.server.ssh.execute('df -t btrfs %s' % self.path)[0]
            self._backingstore = '' if code else 'btrfs'
        return self._backingstore

    def prepare(self):
        if self.network:
            dhcp_start = str(self.network.network + 2)
            dhcp_end = str(self.network.network + self.network.size - 2)
            dhcp_range = ','.join([dhcp_start, dhcp_end])
            values = {
                'USE_LXC_BRIDGE': "true",
                'LXC_BRIDGE': self.config.get('lxc_bridge', 'lxcbr0'),
                'LXC_ADDR': self.network.network + 1,
                'LXC_NETMASK': self.network.netmask,
                'LXC_NETWORK': self.network,
                'LXC_DHCP_RANGE': dhcp_range,
                'LXC_DHCP_MAX': self.network.size - 3,
            }
            config = moves.StringIO()
            for name, value in six.iteritems(values):
                config.write('%(name)s="%(value)s"\n' % {'name': name,
                                                         'value': value})
            config.seek(0)
            self.server.ssh.run('cat > /tmp/.lxc_default', stdin=config)

        self.server.ssh.run('/bin/sh', stdin=_get_script('lxc-install.sh'))
        self.create_local_tunnels()
        self.create_remote_tunnels()

    def create_local_tunnels(self):
        """Create tunel on lxc host side."""
        for tunnel_to in self.config['tunnel_to']:
            script = _get_script_from_template('tunnel-local.sh',
                                               net=self.network,
                                               local=self.server.host,
                                               remote=tunnel_to)
            self.server.ssh.run('/bin/sh', stdin=script)

    def create_remote_tunnels(self):
        """Create tunel on remote side."""
        for tunnel_to in self.config['tunnel_to']:
            script = _get_script_from_template('tunnel-remote.sh',
                                               net=self.network,
                                               local=tunnel_to,
                                               remote=self.server.host)
            server = self._get_updated_server(host=tunnel_to)
            server.ssh.run('/bin/sh', stdin=script)

    def delete_tunnels(self):
        for tunnel_to in self.config['tunnel_to']:
            remote_server = self._get_updated_server(host=tunnel_to)
            remote_server.ssh.execute('ip tun del t%s' % self.network.ip)
            self.server.ssh.execute('ip tun del t%s' % tunnel_to)

    def get_ip(self, name):
        """Get container's ip by name."""
        cmd = 'lxc-attach -n %s ip addr list dev eth0' % name
        for attempt in range(1, 16):
            code, stdout = self.server.ssh.execute(cmd)[:2]
            if code:
                continue
            for line in stdout.splitlines():
                m = INET_ADDR_RE.match(line)
                if m:
                    return m.group(1)
            time.sleep(attempt)
        msg = _('Timeout waiting for ip address of container "%s"') % name
        raise exceptions.TimeoutException(msg)

    def get_port(self, ip):
        """Get forwarded ssh port for instance ip.

        Ssh port forwarding is used for containers access from outside.
        Any container is accessible by host's ip and forwarded port. E.g:

         6.6.6.6:10023 -> 10.1.1.11:22
         6.6.6.6:10024 -> 10.1.1.12:22
         6.6.6.6:10025 -> 10.1.1.13:22

        where 6.6.6.6 is host's ip.

        Ip->port association is stored in self._port_cache to reduce number
        of iptables calls.
        """
        if not self._port_cache:
            self._port_cache = {}
            port_re = re.compile(r'.+ tcp dpt:(\d+).*to:([\d\.]+)\:22')
            cmd = "iptables -n -t nat -L PREROUTING"
            code, out, err = self.server.ssh.execute(cmd)
            for l in out:
                m = port_re.match(l)
                if m:
                    self._port_cache[m.group(2)] = int(m.group(1))
        port = self._port_cache.get(ip)

        if port is None:
            if self._port_cache:
                port = max(self._port_cache.values()) + 1
            else:
                port = 1222
        self._port_cache[ip] = port
        cmd = IPT_PORT_TEMPLATE.format(host_ip=self.server.host, ip=ip,
                                       port=port, action='I')
        self.server.ssh.run(cmd)
        return port

    def create_container(self, name, distribution, release=None):
        cmd = ['lxc-create']
        if self.backingstore == 'btrfs':
            cmd += ['-B', 'btrfs']
        cmd += ['-n', name, '-t', distribution]
        if release:
            if distribution == 'ubuntu':
                cmd += ['--', '-r', release]
            elif distribution == 'debian':
                cmd = ['SUITE=%s' % release] + cmd
        self.server.ssh.run(' '.join(cmd))
        self.configure_container(name)
        self.containers.append(name)

    def create_clone(self, name, source):
        cmd = ['lxc-clone']

        if self.backingstore == 'btrfs':
            cmd.append('--snapshot')
        cmd.extend(['-o', source, '-n', name])
        self.server.ssh.execute(' '.join(cmd))
        self.configure_container(name)
        self.containers.append(name)

    def configure_container(self, name):
        path = os.path.join(self.path, name, 'rootfs')
        conf_script = _get_script('configure_container.sh')
        self.server.ssh.run('/bin/sh -e -s %s' % path, stdin=conf_script)

    def start_containers(self):
        for name in self.containers:
            self.server.ssh.run('lxc-start -d -n %s' % name)

    def stop_containers(self):
        for name in self.containers:
            self.server.ssh.run('lxc-stop -n %s' % name)

    def destroy_ports(self, ipports):
        script = ""
        for ip, port in ipports:
            cmd = IPT_PORT_TEMPLATE.format(action='D', port=port, ip=ip,
                                           host_ip=self.server.host)
            script += cmd + '\n'
        self.server.ssh.run('/bin/sh -e', stdin=script)

    def destroy_containers(self):
        for name in self.containers:
            self.server.ssh.run('lxc-stop -n %s' % name)
            self.server.ssh.run('lxc-destroy -n %s' % name)

    def get_server_object(self, name, wait=True):
        """Create Server object for container."""
        ip = self.get_ip(name)
        if self.config.get('forward_ssh', False):
            server = self._get_updated_server(port=self.get_port(ip))
        else:
            server = self._get_updated_server(host=ip)
        if wait:
            server.ssh.wait(timeout=300)
        return server

    def get_server_objects(self, wait=True):
        """Generate Server objects from all containers."""
        for name in self.containers:
            yield self.get_server_object(name, wait)


class LxcProvider(provider.ProviderFactory):
    """Provide lxc container(s) on given host.

    Sample configuration:
    {
        "type": "LxcProvider",
        "distribution": "ubuntu",
        "start_lxc_network": "10.1.1.0/24",
        "containers_per_host": 32,
        "tunnel_to": ["10.10.10.10"],
        "forward_ssh": false,
        "container_name_prefix": "rally-multinode-02",
        "host_provider": {
            "type": "ExistingServers",
            "credentials": [{"user": "root", "host": "host.net"}]
        }
    }

    """

    CONFIG_SCHEMA = {
        'type': 'object',
        'properties': {
            'type': {'type': 'string'},
            'distribution': {'type': 'string'},
            'release': {'type': 'string'},
            'start_lxc_network': {'type': 'string',
                                  'pattern': '^(\d+\.){3}\d+\/\d+$'},
            'containers_per_host': {'type': 'integer'},
            'forward_ssh': {'type': 'boolean'},
            'tunnel_to': {'type': 'array',
                          'elements': {'type': 'string',
                                       'pattern': '^(\d+\.){3}\d+$'}},
            'container_name_prefix': {'type': 'string'},
            'host_provider': {'type': 'object',
                              'properties': {'type': {'type': 'string'}}},
        },
        'required': ['type', 'containers_per_host',
                     'container_name_prefix', 'host_provider'],

    }

    def validate(self):
        super(LxcProvider, self).validate()
        if 'start_lxc_network' not in self.config:
            return
        lxc_net = netaddr.IPNetwork(self.config['start_lxc_network'])
        num_containers = self.config['containers_per_host']
        if lxc_net.size - 3 < num_containers:
            message = _("Network size is not enough for %d hosts.")
            raise exceptions.InvalidConfigException(message % num_containers)

    def get_host_provider(self):
        return provider.ProviderFactory.get_provider(
            self.config['host_provider'], self.deployment)

    @utils.log_deploy_wrapper(LOG.info, _("Create containers on host"))
    def create_servers(self):
        host_provider = self.get_host_provider()
        name_prefix = self.config['container_name_prefix']
        hosts = []
        if 'start_lxc_network' in self.config:
            network = netaddr.IPNetwork(self.config['start_lxc_network'])
        else:
            network = None
        distribution = self.config.get('distribution', 'ubuntu')
        release = self.config.get('release', None)

        for server in host_provider.create_servers():
            config = {'tunnel_to': self.config.get('tunnel_to', []),
                      'forward_ssh': self.config.get('forward_ssh', False)}
            if network:
                config['network'] = str(network)
            host = LxcHost(server, config)
            host.prepare()
            ip = str(network.ip).replace('.', '-') if network else '0'
            first_name = '%s-000-%s' % (name_prefix, ip)

            host.create_container(first_name, distribution, release)
            for i in range(1, self.config.get('containers_per_host', 1)):
                name = '%s-%03d-%s' % (name_prefix, i, ip)
                host.create_clone(name, first_name)
            host.start_containers()
            hosts.append(host)

            if network:
                network += 1

        servers = []

        for host in hosts:
            for server in host.get_server_objects():
                servers.append(server)
            info = {'host': host.server.get_credentials(),
                    'config': host.config,
                    'forwarded_ports': host._port_cache.items(),
                    'container_names': host.containers}
            self.resources.create(info)
        return servers

    @utils.log_deploy_wrapper(LOG.info, _("Destroy host(s)"))
    def destroy_servers(self):
        for resource in self.resources.get_all():
            server = provider.Server.from_credentials(resource['info']['host'])
            lxc_host = LxcHost(server, resource['info']['config'])
            lxc_host.containers = resource['info']['container_names']
            lxc_host.destroy_containers()
            lxc_host.destroy_ports(resource['info']['forwarded_ports'])
            lxc_host.delete_tunnels()
            self.resources.delete(resource['id'])
        host_provider = self.get_host_provider()
        host_provider.destroy_servers()
