# coding=utf-8
"""Amazon Web Services EC2

Security: IAM Role and policy:
    By default, the EC2 instance is configured to use a IAM role and policy
    generated by Apyfal on first instantiation.

    This give access in read and write to all S3 Buckets.
    This behavior can be changed by overriding IAM Role and/or policy.

    This can be done using by:
     - Selecting an already existing role or policy on class instantiation
       (using ``role`` or ``policy`` argument).
     - Updating policy and role documents in class attributes before
       instantiating the class for the first time (``POLICY_DOCUMENT`` and
       ``ASSUME_ROLE_POLICY_DOCUMENT`` class attributes).
     - Modifying policy or role parameters generated by Apyfal in AIM console.

    Notes that FPGA related rules in policy are mandatory to any Accelerator use
    and needs to be present in custom policy.

    Apyfal create the default AIM rules and policy only if missing and will
    not overwrite existing.
"""

from concurrent.futures import (ThreadPoolExecutor as _ThreadPoolExecutor,
                                as_completed as _as_completed)
from contextlib import contextmanager as _contextmanager
from json import dumps as _json_dumps

import boto3 as _boto3
from botocore.exceptions import ClientError as _ClientError

from apyfal.host._csp import CSPHost as _CSPHost
import apyfal.exceptions as _exc
import apyfal._utilities as _utl
from apyfal._utilities import get_logger as _get_logger


@_contextmanager
def _exception_handler(
        to_catch=_ClientError, to_raise=None,
        filter_error_codes=None, exception_msg=None, **exc_kwargs):
    """
    Context manager that catch AWS EC2 exceptions and raises
    Apyfal exceptions.

    Args:
        to_catch (Exception sub class or tuple of Exception): Exception to catch.
            ClientError if not specified.
        to_raise (apyfal.exception.AcceleratorException subclass):
            Exception to raise. apyfal.exceptions.HostRuntimeException if not
            specified.
        filter_error_codes (str or tuple of str):
            Don't raise exception if error code in this argument.
        exception_msg (str): Exception message.
        exc_kwargs: Exception to raise arguments.
    """
    # Performs operation
    try:
        yield

    # Catch specified exceptions
    except to_catch as exception:
        # Try to get error code and message
        try:
            error_dict = exception.response['Error']
            error_code = error_dict['Code']
        except (AttributeError, KeyError):
            raise _exc.HostRuntimeException(exception_msg, exc=exception)

        # Converts single str to tuple
        if filter_error_codes is None:
            filter_error_codes = ()
        elif isinstance(filter_error_codes, str):
            filter_error_codes = (filter_error_codes,)

        # Raises if not in filter
        if error_code not in filter_error_codes:
            exception = to_raise or _exc.HostRuntimeException
            raise exception(
                exception_msg, exc=error_dict['Message'], **exc_kwargs)


class AWSHost(_CSPHost):
    """AWS EC2 CSP

    Args:
        host_type (str): Cloud service provider name. Default to "AWS".
        config (apyfal.configuration.Configuration, path-like object or file-like object):
            If not set, will search it in current working directory,
            in current user "home" folder. If none found, will use default
            configuration values.
            Path-like object can be path, URL or cloud object URL.
        client_id (str): AWS Access Key ID.
        secret_id (str): AWS Secret Access Key.
        region (str): AWS region. Needs a EC2 region supporting instances with
            FPGA devices.
        instance_type (str): AWS EC2 Instance type.
            Default defined by accelerator.
        key_pair (str): AWS Key pair. Default to 'AccelizeAWSKeyPair'.
        security_group: AWS Security group. Default to 'AccelizeSecurityGroup'.
        instance_id (str): Instance ID of an already existing AWS EC2 instance
            to use. If not specified, create a new instance.
        host_name_prefix (str): Prefix to add to instance name.
            Also used as value of the "Apyfal" tag.
        host_ip (str): IP or URL address of an already existing AWS EC2 instance
            to use. If not specified, create a new instance.
        use_private_ip (bool): If True, on new instances,
            uses private IP instead of public IP as default host IP.
        role (str): AWS IAM role. Generated to allow instance to load AGFI
            (FPGA bitstream) and access to S3. Default to 'AccelizeRole'.
        policy (str): AWS IAM policy. Generated to allow instance to load AGFI
            (FPGA bitstream) and access to S3. Default to 'AccelizePolicy'.
        stop_mode (str or int): Define the "stop" method behavior.
            Default to 'term' if new instance, or 'keep' if already existing
            instance. See "stop_mode" property for more information and possible
            values.
        init_config (bool or apyfal.configuration.Configuration, path-like object or file-like object):
            Configuration file to pass to instance on
            initialization. This configuration file will be used as default for
            host side accelerator.
            If value is True, use 'config' configuration.
            If value is a configuration use this configuration.
            If value is None or False, don't passe any configuration file
            (This is default behavior).
        init_script (path-like object or file-like object): A bash script
            to execute on instance startup.
        ssl_cert_crt (path-like object or file-like object or bool):
            Public ".crt" key file of the SSL ssl_cert_key used to provides
            HTTPS.
            If not specified, uses already generated certificate if found.
            If False, disable HTTPS.
        ssl_cert_key (path-like object or file-like object):
            Private ".key" key file of the SSL ssl_cert_key used to provides
            HTTPS.
            If not specified, uses already generated key if found.
        ssl_cert_generate (bool): Generate a self signed ssl_cert_key.
            The ssl_cert_key and private key will be stored in files specified
            by "ssl_cert_crt" and "ssl_cert_key" (Or temporary certificates if
            not specified). Note that this ssl_cert_key is only safe if other
            client verify it by providing "ssl_cert_crt". No Certificate
            Authority are available to trust this ssl_cert_key.
        delete_volumes_on_termination (bool): If True, force the volumes
            deletion on instance termination, else use AWS default behavior
            (Not deleting volumes except root device). Default to True.
        spot_instance (bool): If True, use spot instance. Default to True.
        spot_block_duration (int): Spot instance blocks duration in minutes.
            Must be multiple of 60, >= 60 and <= 360.
            If set to None, use AWS default spot block (60 minutes).
        boto3_session_kwargs (dict): Extra keyword arguments for
            boto3.session.Session
        boto3_client_kwargs (dict): Extra keyword arguments for
            boto3.session.Session.client.
        boto3_create_instances_kwargs (dict): Extra Keyword arguments for
            boto3.session.Session.resource('ec2').create_instances.
    """
    #: Provider name to use
    NAME = 'AWS'

    #: AWS Website
    DOC_URL = "https://aws.amazon.com"

    #: IAM Policy document: Grant FPGA Access and S3 Buckets access
    POLICY_DOCUMENT = {"Version": "2012-10-17", "Statement": [

        # Grant FPGA access
        {"Sid": "AllowFpgaCommands", "Effect": "Allow",
         "Action": ["ec2:AssociateFpgaImage", "ec2:DisassociateFpgaImage",
                    "ec2:DescribeFpgaImages"], "Resource": ["*"]},

        # Grant S3 buckets access
        {"Sid": "AllowS3Access", "Effect": "Allow",
         "Action": ["s3:GetObject", "s3:PutObject", "s3:ListBucket",
                    "s3:ListAllMyBuckets"], "Resource": ["arn:aws:s3:::*"]}]}

    #: IAM Assume Role Policy Document: Grant rights defined by policy to role.
    ASSUME_ROLE_POLICY_DOCUMENT = {
        "Version": "2012-10-17", "Statement": {
            "Effect": "Allow", "Principal": {"Service": "ec2.amazonaws.com"},
            "Action": "sts:AssumeRole"}}

    #: Spot instance configuration
    SPOT_INSTANCE_MARKET_OPTIONS = {
        'InstanceMarketOptions': {
            'MarketType': 'spot',
            'SpotOptions': {
                'SpotInstanceType': 'one-time',
                'InstanceInterruptionBehavior': 'terminate',
            }
        }
    }

    STATUS_RUNNING = 'running'
    STATUS_STOPPED = 'stopped'
    STATUS_STOPPING = 'stopping'

    _INFO_NAMES = _CSPHost._INFO_NAMES.copy()
    _INFO_NAMES.update(['_role', '_policy'])

    def __init__(self, role=None, policy=None,
                 delete_volumes_on_termination=None, spot_instance=None,
                 spot_block_duration=None,  boto3_session_kwargs=None,
                 boto3_client_kwargs=None, boto3_create_instances_kwargs=None,
                 **kwargs):
        _CSPHost.__init__(self, **kwargs)

        # Get AWS specific arguments
        section = self._config[self._config_section]
        self._role, self._policy = self._get_role_and_policy(role, policy)
        self._policy_arn = None
        self._instance_profile_name = 'AccelizeLoadFPGA'
        self._delete_volumes_on_termination = (
            delete_volumes_on_termination or
            section.get_literal('delete_volumes_on_termination') or True)
        self._spot_instance = (
                spot_instance or section.get_literal('spot_instance') or True)
        self._spot_block_duration = (
                spot_block_duration or section.get_literal(
            'spot_block_duration'))
        self._block_devices = None

        # Session, clients and resources are lazy instantiated
        self._boto3_session_kwargs = (
            boto3_session_kwargs or
            section.get_literal('boto3_session_kwargs') or dict())
        self._boto3_client_kwargs = (
            boto3_client_kwargs or
            section.get_literal('boto3_client_kwargs') or dict())
        self._boto3_create_instances_kwargs = (
            boto3_create_instances_kwargs or
            section.get_literal('boto3_create_instances_kwargs') or dict())

    @property
    @_utl.memoizedmethod
    def _session(self):
        """
        Return boto3 session

        Returns:
            boto3.session.Session: session
        """
        kwargs = dict(
            aws_access_key_id=self._client_id,
            aws_secret_access_key=self._secret_id, region_name=self._region)
        kwargs.update(self._boto3_session_kwargs)
        return _boto3.session.Session(**kwargs)

    @property
    @_utl.memoizedmethod
    def _ec2_client(self):
        """
        Return boto3 EC2 client

        Returns:
            boto3.session.Session.client: client
        """
        return self._session.client('ec2', **self._boto3_client_kwargs)

    @property
    @_utl.memoizedmethod
    def _ec2_resource(self):
        """
        Return boto3 EC2 resource

        Returns:
            boto3.session.Session.resource: resource
        """
        return self._session.resource('ec2')

    @property
    @_utl.memoizedmethod
    def _iam_client(self):
        """
        Return boto3 IAM client

        Returns:
            boto3.session.Session.client: client
        """
        return self._session.client('iam', **self._boto3_client_kwargs)

    @property
    @_utl.memoizedmethod
    def _iam_resource(self):
        """
        Return boto3 IAM resource

        Returns:
            boto3.session.Session.resource: resource
        """
        return self._session.resource('iam')

    def _check_credential(self):
        """
        Check CSP credentials.

        Raises:
            apyfal.exceptions.HostAuthenticationException:
                Authentication failed.
        """
        with _exception_handler(to_raise=_exc.HostAuthenticationException):
            self._ec2_client.describe_key_pairs()

    def _init_key_pair(self):
        """
        Initializes key pair.
        """
        # Checks if Key pairs exists, needs to get the full pairs list
        # and compare in lower case because Boto perform its checks case
        # sensitive and AWS use case insensitive names.
        with _exception_handler():
            key_pairs = self._ec2_client.describe_key_pairs()

        name_lower = self._key_pair.lower()
        for key_pair in key_pairs['KeyPairs']:
            key_pair_name = key_pair['KeyName']
            if key_pair_name.lower() == name_lower:
                self._key_pair = key_pair_name
                return

        # Key does not exist on the CSP, create it
        with _exception_handler():
            key_pair = self._ec2_resource.create_key_pair(
                KeyName=self._key_pair)

        _utl.create_key_pair_file(self._key_pair, key_pair.key_material)
        _get_logger().debug(_utl.gen_msg(
            "created_named", "key pair", self._key_pair))

    def _init_policy(self):
        """
        Initialize IAM policy.

        This policy allow instance to:
            - Load FPGA bitstream.
            - Access to S3 buckets objects for read and write.

        Returns:
            str: 'policy'
        """
        # Create a policy
        with _exception_handler(filter_error_codes='EntityAlreadyExists'):
            self._iam_client.create_policy(
                PolicyName=self._policy,
                PolicyDocument=_json_dumps(self.POLICY_DOCUMENT))

            _get_logger().debug(_utl.gen_msg(
                'created_named', 'policy', self._policy))

        with _exception_handler():
            response = self._iam_client.list_policies(
                Scope='Local', OnlyAttached=False, MaxItems=100)
        for policy_item in response['Policies']:
            if policy_item['PolicyName'] == self._policy:
                self._policy_arn = policy_item['Arn']
                # 'policy' returns str is used set future object ID
                return 'policy'

        raise _exc.HostConfigurationException(
            gen_msg=('created_failed_named', 'IAM policy', self._policy))

    def _init_role(self):
        """
        Initialize IAM role.

        This role allow to perform actions defined by policy.
        """
        with _exception_handler(filter_error_codes='EntityAlreadyExists'):
            role = self._iam_resource.create_role(
                RoleName=self._role,
                AssumeRolePolicyDocument=_json_dumps(
                    self.ASSUME_ROLE_POLICY_DOCUMENT),
                Description=_utl.gen_msg('accelize_generated'))

            _get_logger().debug(_utl.gen_msg('created_named', 'IAM role', role))

    def _attach_role_policy(self):
        """
        Attach IAM policy to IAM role.
        """
        with _exception_handler(filter_error_codes='EntityAlreadyExists'):
            self._iam_client.attach_role_policy(
                PolicyArn=self._policy_arn, RoleName=self._role)

            _get_logger().debug(_utl.gen_msg(
                'attached_to', 'IAM policy', self._policy_arn,
                'IAM role', self._role))

    def _init_instance_profile(self):
        """
        Initialize IAM instance profile.

        This instance_profile allow to perform actions defined by role.
        """
        # Create instance profile
        with _exception_handler(filter_error_codes='EntityAlreadyExists'):
            self._iam_client.create_instance_profile(
                InstanceProfileName=self._instance_profile_name)

            _get_logger().debug(_utl.gen_msg(
                'created_named', 'IAM instance profile',
                self._instance_profile_name))

    def _attach_instance_profile_role(self):
        """
        Attach IAM role to IAM instance profile.
        """
        # Attach role to instance profile
        with _exception_handler(filter_error_codes='LimitExceeded'):
            with _utl.Timeout(self.TIMEOUT,
                              sleep=self._TIMEOUT_SLEEP) as timeout:
                while True:
                    try:
                        self._iam_client.add_role_to_instance_profile(
                            InstanceProfileName=self._instance_profile_name,
                            RoleName=self._role)
                        break

                    # Some time, instance_profile is not ready immediately
                    except _ClientError as exception:
                        if (exception.response['Error']['Code'] ==
                                'NoSuchEntityException'):
                            if timeout.reached():
                                raise _exc.HostRuntimeException(gen_msg=(
                                    'timeout',
                                    'IAM instance_profile and role attachment'))
                            continue
                        raise

            _get_logger().debug(_utl.gen_msg(
                'attached_to', 'IAM role', self._role, 'IAM instance profile',
                self._instance_profile_name))

    def _init_security_group(self):
        """
        Initialize security group.
        """
        # Get list of security groups
        # Checks if Key pairs exists, like for key pairs
        # needs  case insensitive names check
        with _exception_handler():
            security_groups = self._ec2_client.describe_security_groups()

        name_lower = self._security_group.lower()
        group_exists = False
        security_group_id = ''
        for security_group in security_groups['SecurityGroups']:
            group_name = security_group['GroupName']
            if group_name.lower() == name_lower:
                # Update name
                self._security_group = group_name

                # Get group ID
                security_group_id = security_group['GroupId']

                # Mark as existing
                group_exists = True
                break

        # Try to create security group if not exist
        if not group_exists:
            # Get VPC
            with _exception_handler():
                vpc_id = self._ec2_client.describe_vpcs().get(
                    'Vpcs', [{}])[0].get('VpcId', '')

            with _exception_handler():
                response = self._ec2_client.create_security_group(
                    GroupName=self._security_group,
                    Description=_utl.gen_msg('accelize_generated'),
                    VpcId=vpc_id)

            # Get group ID
            security_group_id = response['GroupId']

            _get_logger().debug(_utl.gen_msg(
                'created_named', 'security group', security_group_id))

        # Add host IP to security group if not already done
        public_ip = _utl.get_host_public_ip()

        ip_permissions = []
        for port in self.ALLOW_PORTS:
            ip_permissions.append({
                'IpProtocol': 'tcp', 'FromPort': port, 'ToPort': port,
                'IpRanges': [{'CidrIp': public_ip}]})

        with _exception_handler(
                filter_error_codes='InvalidPermission.Duplicate'):
            self._ec2_client.authorize_security_group_ingress(
                GroupId=security_group_id, IpPermissions=ip_permissions)

        _get_logger().debug(
            _utl.gen_msg('authorized_ip', public_ip, self._security_group))

    def _init_block_device_mappings(self):
        """
        Initialize block device mapping
        """
        if self._delete_volumes_on_termination:
            # Get block devices for specified image
            with _exception_handler():
                block_devices = self._ec2_resource.Image(
                    self._image_id).block_device_mappings

            # Configure block devices
            for block_device in block_devices:
                ebs = block_device.get('Ebs')
                if ebs is None:
                    continue

                # AWS don't allow both SnapshotId and Encrypted values
                if 'SnapshotId' in ebs and 'Encrypted' in ebs:
                    del ebs['Encrypted']

                # Force termination if not forced
                if not ebs.get('DeleteOnTermination', False):
                    ebs['DeleteOnTermination'] = True

                    _get_logger().debug(
                        'Enable "DeleteOnTermination" on "%s" volume.',
                        block_device['DeviceName'])

                    # Pass parameter if at least a disk parameter changed
                    self._block_devices = block_devices

    def _get_instance(self):
        """
        Returns current instance.

        Returns:
            object: Instance
        """
        with _exception_handler(gen_msg=('no_instance_id', self._instance_id)):
            return self._ec2_resource.Instance(self._instance_id)

    def _get_public_ip(self):
        """
        Read current instance public IP from CSP instance.

        Returns:
            str: IP address
        """
        with _exception_handler(gen_msg='no_instance_ip'):
            return self._instance.public_ip_address

    def _get_private_ip(self):
        """
        Read current instance private IP from CSP instance.

        Returns:
            str: IP address
        """
        with _exception_handler(gen_msg='no_instance_ip'):
            return self._instance.private_ip_address

    def _get_status(self):
        """
        Returns current status of current instance.

        Returns:
            str: Status
        """
        with _utl.Timeout(1, sleep=0.01) as timeout:
            while True:
                # Check Timeout
                if timeout.reached():
                    raise _exc.HostRuntimeException(gen_msg=(
                        'no_instance_id', self._instance_id))

                # Get status
                with _exception_handler(
                        filter_error_codes='InvalidInstanceID.NotFound'):
                    return self._instance.state["Name"]

    def _create_instance(self):
        """
        Initialize and create instance.
        """
        futures = []
        with _ThreadPoolExecutor(max_workers=6) as executor:
            # Run configuration in parallel
            policy = executor.submit(self._init_policy)
            role = executor.submit(self._init_role)
            instance_profile = executor.submit(self._init_instance_profile)
            for method in (self._init_key_pair, self._init_security_group,
                           self._init_block_device_mappings):
                futures.append(executor.submit(method))

            # Wait that role, instance_profile and policy are completed
            # to attach them
            for future in _as_completed((policy, instance_profile)):
                role.result()
                futures.append(executor.submit(
                    self._attach_role_policy if future.result() == 'policy' else
                    self._attach_instance_profile_role))

        # Wait completion
        for future in _as_completed(futures):
            future.result()

    def _start_new_instance(self):
        """
        Start a new instance.

        Returns:
            object: Instance
            str: Instance ID
        """
        kwargs = dict(
                ImageId=self._image_id, InstanceType=self._instance_type,
                KeyName=self._key_pair, SecurityGroups=[self._security_group],
                IamInstanceProfile={'Name': 'AccelizeLoadFPGA'},
                InstanceInitiatedShutdownBehavior='terminate',
                TagSpecifications=[{'ResourceType': 'instance', 'Tags': [
                    {'Key': 'Generated',
                     'Value': _utl.gen_msg('accelize_generated')},
                    {'Key': 'Name', 'Value': self._get_host_name()},
                    {'Key': 'Apyfal', 'Value': self._get_tag()}]}],
                MinCount=1, MaxCount=1, UserData=self._user_data)

        if self._block_devices:
            kwargs['BlockDeviceMappings'] = self._block_devices

        if self._spot_instance:
            _utl.recursive_update(kwargs, self.SPOT_INSTANCE_MARKET_OPTIONS)
            if self._spot_block_duration:
                kwargs['InstanceMarketOptions']['SpotOptions'][
                    'BlockDurationMinutes'] = self._spot_block_duration

        _utl.recursive_update(kwargs, self._boto3_create_instances_kwargs)

        # Create instance
        with _exception_handler():
            instance = self._ec2_resource.create_instances(**kwargs)[0]

        return instance, instance.id

    def _start_existing_instance(self, status):
        """
        Start a existing instance.

        Args:
            status (str): Status of the instance.
        """
        # Waiting for the instance stop if currently stopping
        if status == self.STATUS_STOPPING:
            with _utl.Timeout(self.TIMEOUT,
                              sleep=self._TIMEOUT_SLEEP) as timeout:
                while True:
                    # Get instance status
                    status = self._status()
                    if status != self.STATUS_STOPPING:
                        break
                    elif timeout.reached():
                        raise _exc.HostRuntimeException(gen_msg=(
                            'timeout_status', 'stop', status))

        # If instance stopped, starts it
        if status == self.STATUS_STOPPED:
            with _exception_handler():
                self._instance.start()

        # If another status, raises error
        elif status != self.STATUS_RUNNING:
            raise _exc.HostRuntimeException(gen_msg=(
                'unable_to_status', 'start', status))

    def _terminate_instance(self):
        """
        Terminate and delete instance.
        """
        if self._instance is not None:
            with _exception_handler():
                return self._instance.terminate()

    def _pause_instance(self):
        """
        Pause instance.
        """
        if self._instance is not None:
            with _exception_handler():
                return self._instance.stop()

    def _iter_hosts(self):
        """
        Iterates over accelerator hosts of current type.

        Returns:
            generator of dict: dicts contains attributes values of the host.
        """
        with _exception_handler():
            for instance in self._ec2_resource.instances.filter(
                    Filters=[{'Name': 'instance-state-name',
                              'Values': ['running']}]):
                host_name = [instance.tags[index]['Value']
                             for index in range(len(instance.tags))
                             if instance.tags[index]['Key'] == 'Name'][0]

                # Yields only matching accelerator instances
                if self._is_accelerator_host(host_name):
                    yield dict(
                        instance_id=instance.id,
                        instance_type=instance.instance_type,
                        private_ip=instance.private_ip_address,
                        public_ip=instance.public_ip_address,
                        host_name=host_name,
                        security_group=instance.security_groups[0]['GroupName'],
                        image_id=instance.image_id,
                        key_pair=instance.key_name)
