#!/usr/bin/env python
#
# Copyright 2014 (c) Lei Xu <eddyxu@gmail.com>
#
# 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.

"""
The basic building block of running VSFS and benchmarks on EC2
"""

from __future__ import print_function
from datetime import datetime
from fabric.api import cd, env, sudo, settings, put, run, execute, parallel
from fabric.colors import yellow
import boto.ec2
import boto.vpc
import fabric.exceptions
import yaml
import glob
import logging
import os
import time

INSTALL_SCRIPT = \
    'https://raw.github.com/vsfs/vsfs-devtools/master/install-devbox.sh'
VSFS_GIT_URL = 'https://github.com/vsfs/vsfs.git'


class Volume(object):
    """Encaputure operations for EBS volume.
    """
    def __init__(self, ec2, volume=None, volume_id=None):
        if not volume and not volume_id:
            raise RuntimeError("You must specify one of Volume and VolumeId")
        self.ec2 = ec2
        if volume:
            self.vol = volume
        else:
            volumes = self.ec2.get_all_volumes(volume_ids=[volume_id])
            assert len(volumes) == 1
            self.vol = volumes[0]

    def wait(self, state, sleep_time=10):
        while True:
            logging.info('The volume (%s) state is %s..', self.vol.id,
                         self.vol.attachment_state())
            if self.vol.attachment_state() == state:
                break
            print(yellow('Sleep %d seconds...' % sleep_time))
            time.sleep(sleep_time)
            self.vol.update()

    def detach(self, force=False):
        if self.attached():
            self.vol.detach(force=force)
            self.wait(None)
        return True

    def status(self):
        return self.vol.state()

    def attached(self):
        return self.vol.attachment_state() == 'attached'


class SpotRequest(object):
    """A Spot Instance request
    """
  
  def __init__(self, ec2, ami, count=1, subnet=None,
                 security_group_ids=None):
        """Initialize a spot instance reques
t with ec2 connection and AMI.

        @param conn an established EC2 connection.
        @param ami an AMI id.

        Optional parameters:
        @param count the number of instances requested.
        """
        self.ec2 = ec2
        self.conn = self.ec2.conn
        self.ami = ami
        self.security_groups = self.ec2.security_groups
        self.key_name = self.ec2.key_name
        self.count = count
        self.requests = None
        self.instances = []
        # For VPC
        self.subnet = subnet
        self.security_group_ids = security_group_ids

        if self.subnet and not self.security_group_ids:
            message = 'Issuing an Spot Instance Request in VPC requires ' \
                      'security_group_ids being set.'
            logging.fatal(message)
            raise RuntimeError(message)

    @staticmethod
    def to_string(request):
        """Format a boto.ec2.SpotInstanceRequest to string.
        """
        return 'Spot instance request: request_id={}, price={}, ' \
               'instance_id={}, type={}, state={}'.format(
                   request.id, request.price, request.instance_id,
                   request.type, request.state)

    @staticmethod
    def print_all_requests(conn):
        """Prints infomation of all requests.

        @param conn an established boto.EC2Connection.
        """
        logging.info('Querying for all spot instance requests..')
        requests = conn.get_all_spot_instance_requests()
        for request in requests:
            if request.state == 'cancelled':
                continue
            print(SpotRequest.to_string(request))

    def __str__(self):
        """Return string representation of a spot request.
        """
        return SpotRequest.to_string(self.request)

    def __del__(self):
        """Cancel all left-over spot instance requests.
        """
        self.cancel()

    def request(self, price, instance_type='m1.small'):
        """Issue an new Spot Instance request.
        """
        if self.subnet:  # Request within VPC
            self.requests = self.conn.request_spot_instances(
                price, self.ami, count=self.count, instance_type=instance_type,
                subnet_id=self.subnet,
                security_group_ids=self.security_group_ids,
                key_name=self.key_name)
        else:  # Normal Spot Requests.
            self.requests = self.conn.request_spot_instances(
                price, self.ami, count=self.count, instance_type=instance_type,
                security_groups=self.security_groups,
                key_name=self.key_name)

    def cancel(self):
        """Cancels the request.
        """
        if self.requests:
            for request in self.requests:
                request.cancel()

    def ready(self, no_wait=False, sleep_time=30):
        """Returns true when the instances are ready.
        """
        if not self.request:
            logging.error('No request has been issued.')
            return False
        try:
            while True:
                requests = self.conn.get_all_spot_instance_requests(
                    request_ids=[r.id for r in self.requests])
                if len(requests) == 0:
                    print(yellow('Sleep %d seconds...' % sleep_time))
                    print(yellow('Press Ctrl-C to cancel...'))
                    if no_wait:
                        return False
                    time.sleep(sleep_time)
                    continue
                succ_instance_ids = []  # success instances
                for request in requests:
                    logging.info(SpotRequest.to_string(request))
                    if request.state == 'cancelled':
                        logging.info('The spot request has been canceled.')
                        logging.info(SpotRequest.to_string(request))
                        return False
                    elif request.state == 'active':
                        logging.info('Got instances, lets roll..: '
                                     'instance_id=%s', request.instance_id)
                        succ_instance_ids.append(request.instance_id)
                # All instance requests are full filled.
                if len(succ_instance_ids) == len(self.requests):
                    logging.info('All instances are full filled.')
                    self.instances = []
                    for reservation in self.conn.get_all_instances(
                            instance_ids=succ_instance_ids):
                        self.instances.extend(reservation.instances)
                    return True
                if no_wait:
                    return False
                logging.info('Sleep for %d seconds...', sleep_time)
                time.sleep(sleep_time)
        except KeyboardInterrupt:
            print(yellow('Ctrl-C received, canceling the spot request.'))
            self.cancel()


class EC2Instance(object):
    """An EC2 instance.
    """
    USER = 'ubuntu'

    def __init__(self, ec2_instance, key_file):
        self.instance = ec2_instance
        self.key_file = key_file
        self.retry = 10
        self.sleep_time = 30

    @staticmethod
    def wait_running(instance, sleep_time=30):
        """
        """
        wait_for_boot = 1
        while True:
            instance.update()
            if instance.state == 'terminated':
                return False
            if instance.state == 'pending':
                print(yellow('Waitting for booting...{} times'.format(
                             wait_for_boot)))
                wait_for_boot += 1
                time.sleep(sleep_time)
                continue
   