# Copyright 2018 Amazon.com, Inc. or its affiliates. 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. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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 hashlib
import json
import os
import sys

from botocore import compat, config
from botocore.exceptions import ClientError
from awscli.compat import compat_open
from awscli.customizations.ecs import exceptions, filehelpers
from awscli.customizations.commands import BasicCommand

TIMEOUT_BUFFER_MIN = 10
DEFAULT_DELAY_SEC = 15
MAX_WAIT_MIN = 360  # 6 hours


class ECSDeploy(BasicCommand):
    NAME = 'deploy'

    DESCRIPTION = (
        "Deploys a new task definition to the specified ECS service. "
        "Only services that use CodeDeploy for deployments are supported. "
        "This command will register a new task definition, update the "
        "CodeDeploy appspec with the new task definition revision, create a "
        "CodeDeploy deployment, and wait for the deployment to successfully "
        "complete. This command will exit with a return code of 255 if the "
        "deployment does not succeed within 30 minutes by default or "
        "up to 10 minutes more than your deployment group's configured wait "
        "time (max of 6 hours)."
    )

    ARG_TABLE = [
        {
            'name': 'service',
            'help_text': (
                "The short name or full Amazon Resource Name "
                "(ARN) of the service to update"
            ),
            'required': True,
        },
        {
            'name': 'task-definition',
            'help_text': (
                "The file path where your task definition file is "
                "located. The format of the file must be the same "
                "as the JSON output of: <codeblock>aws ecs "
                "register-task-definition "
                "--generate-cli-skeleton</codeblock>"
            ),
            'required': True,
        },
        {
            'name': 'codedeploy-appspec',
            'help_text': (
                "The file path where your AWS CodeDeploy appspec "
                "file is located. The appspec file may be in JSON "
                "or YAML format. The <code>TaskDefinition</code> "
                "property will be updated within the appspec with "
                "the newly registered task definition ARN, "
                "overwriting any placeholder values in the file."
            ),
            'required': True,
        },
        {
            'name': 'cluster',
            'help_text': (
                "The short name or full Amazon Resource Name "
                "(ARN) of the cluster that your service is "
                "running within. If you do not specify a "
                "cluster, the \"default\" cluster is assumed."
            ),
            'required': False,
        },
        {
            'name': 'codedeploy-application',
            'help_text': (
                "The name of the AWS CodeDeploy application "
                "to use for the deployment. The specified "
                "application must use the 'ECS' compute "
                "platform. If you do not specify an "
                "application, the application name "
                "<code>AppECS-[CLUSTER_NAME]-[SERVICE_NAME]</code> "
                "is assumed."
            ),
            'required': False,
        },
        {
            'name': 'codedeploy-deployment-group',
            'help_text': (
                "The name of the AWS CodeDeploy deployment "
                "group to use for the deployment. The "
                "specified deployment group must be associated "
                "with the specified ECS service and cluster. "
                "If you do not specify a deployment group, "
                "the deployment group name "
                "<code>DgpECS-[CLUSTER_NAME]-[SERVICE_NAME]</code> "
                "is assumed."
            ),
            'required': False,
        },
    ]

    MSG_TASK_DEF_REGISTERED = (
        "Successfully registered new ECS task definition {arn}\n"
    )

    MSG_CREATED_DEPLOYMENT = "Successfully created deployment {id}\n"

    MSG_SUCCESS = "Successfully deployed {task_def} to service '{service}'\n"

    USER_AGENT_EXTRA = 'md/customization#ecs-deploy'

    def _run_main(self, parsed_args, parsed_globals):
        register_task_def_kwargs, appspec_obj = self._load_file_args(
            parsed_args.task_definition, parsed_args.codedeploy_appspec
        )

        ecs_client_wrapper = ECSClient(
            self._session, parsed_args, parsed_globals, self.USER_AGENT_EXTRA
        )

        self.resources = self._get_resource_names(
            parsed_args, ecs_client_wrapper
        )

        codedeploy_client = self._session.create_client(
            'codedeploy',
            region_name=parsed_globals.region,
            verify=parsed_globals.verify_ssl,
            config=config.Config(user_agent_extra=self.USER_AGENT_EXTRA),
        )

        self._validate_code_deploy_resources(codedeploy_client)

        self.wait_time = self._cd_validator.get_deployment_wait_time()

        self.task_def_arn = self._register_task_def(
            register_task_def_kwargs, ecs_client_wrapper
        )

        self._create_and_wait_for_deployment(codedeploy_client, appspec_obj)
        return 0

    def _create_and_wait_for_deployment(self, client, appspec):
        deployer = CodeDeployer(client, appspec)
        deployer.update_task_def_arn(self.task_def_arn)
        deployment_id = deployer.create_deployment(
            self.resources['app_name'], self.resources['deployment_group_name']
        )

        sys.stdout.write(self.MSG_CREATED_DEPLOYMENT.format(id=deployment_id))

        deployer.wait_for_deploy_success(deployment_id, self.wait_time)
        service_name = self.resources['service']

        sys.stdout.write(
            self.MSG_SUCCESS.format(
                task_def=self.task_def_arn, service=service_name
            )
        )
        sys.stdout.flush()

    def _get_file_contents(self, file_path):
        full_path = os.path.expandvars(os.path.expanduser(file_path))
        try:
            with compat_open(full_path) as f:
                return f.read()
        except (OSError, IOError, UnicodeDecodeError) as e:
            raise exceptions.FileLoadError(file_path=file_path, error=e)

    def _get_resource_names(self, args, ecs_client):
        service_details = ecs_client.get_service_details()
        service_name = service_details['service_name']
        cluster_name = service_details['cluster_name']

        application_name = filehelpers.get_app_name(
            service_name, cluster_name, args.codedeploy_application
        )
        deployment_group_name = filehelpers.get_deploy_group_name(
            service_name, cluster_name, args.codedeploy_deployment_group
        )

        return {
            'service': service_name,
            'service_arn': service_details['service_arn'],
            'cluster': cluster_name,
            'cluster_arn': service_details['cluster_arn'],
            'app_name': application_name,
            'deployment_group_name': deployment_group_name,
        }

    def _load_file_args(self, task_def_arg, appspec_arg):
        task_def_string = self._get_file_contents(task_def_arg)
        register_task_def_kwargs = json.loads(task_def_string)

        appspec_string = self._get_file_contents(appspec_arg)
        appspec_obj = filehelpers.parse_appspec(appspec_string)

        return register_task_def_kwargs, appspec_obj

    def _register_task_def(self, task_def_kwargs, ecs_client):
        response = ecs_client.register_task_definition(task_def_kwargs)

        task_def_arn = response['taskDefinition']['taskDefinitionArn']

        sys.stdout.write(self.MSG_TASK_DEF_REGISTERED.format(arn=task_def_arn))
        sys.stdout.flush()

        return task_def_arn

    def _validate_code_deploy_resources(self, client):
        validator = CodeDeployValidator(client, self.resources)
        validator.describe_cd_resources()
        validator.validate_all()
        self._cd_validator = validator


class CodeDeployer:
    MSG_WAITING = (
        "Waiting for {deployment_id} to succeed "
        "(will wait up to {wait} minutes)...\n"
    )

    def __init__(self, cd_client, appspec_dict):
        self._client = cd_client
        self._appspec_dict = appspec_dict

    def create_deployment(self, app_name, deploy_grp_name):
        request_obj = self._get_create_deploy_request(
            app_name, deploy_grp_name
        )

        try:
            response = self._client.create_deployment(**request_obj)
        except ClientError as e:
            raise exceptions.ServiceClientError(
                action='create deployment', error=e
            )

        return response['deploymentId']

    def _get_appspec_hash(self):
        appspec_str = json.dumps(self._appspec_dict)
        appspec_encoded = compat.ensure_bytes(appspec_str)
        return hashlib.sha256(appspec_encoded).hexdigest()

    def _get_create_deploy_request(self, app_name, deploy_grp_name):
        return {
            "applicationName": app_name,
            "deploymentGroupName": deploy_grp_name,
            "revision": {
                "revisionType": "AppSpecContent",
                "appSpecContent": {
                    "content": json.dumps(self._appspec_dict),
                    "sha256": self._get_appspec_hash(),
                },
            },
        }

    def update_task_def_arn(self, new_arn):
        """
        Inserts the ARN of the previously created ECS task definition
        into the provided appspec.

        Expected format of ECS appspec (YAML) is:
            version: 0.0
            resources:
              - <service-name>:
                  type: AWS::ECS::Service
                  properties:
                    taskDefinition: <value>  # replace this
                    loadBalancerInfo:
                      containerName: <value>
                      containerPort: <value>
        """
        appspec_obj = self._appspec_dict

        resources_key = filehelpers.find_required_key(
            'codedeploy-appspec', appspec_obj, 'resources'
        )
        updated_resources = []

        # 'resources' is a list of string:obj dictionaries
        for resource in appspec_obj[resources_key]:
            for name in resource:
                # get content of resource
                resource_content = resource[name]
                # get resource properties
                properties_key = filehelpers.find_required_key(
                    name, resource_content, 'properties'
                )
                properties_content = resource_content[properties_key]
                # find task definition property
                task_def_key = filehelpers.find_required_key(
                    properties_key, properties_content, 'taskDefinition'
                )

                # insert new task def ARN into resource
                properties_content[task_def_key] = new_arn

            updated_resources.append(resource)

        appspec_obj[resources_key] = updated_resources
        self._appspec_dict = appspec_obj

    def wait_for_deploy_success(self, id, wait_min):
        waiter = self._client.get_waiter("deployment_successful")

        if wait_min is not None and wait_min > MAX_WAIT_MIN:
            wait_min = MAX_WAIT_MIN

        elif wait_min is None or wait_min < 30:
            wait_min = 30

        delay_sec = DEFAULT_DELAY_SEC
        max_attempts = (wait_min * 60) / delay_sec
        config = {'Delay': delay_sec, 'MaxAttempts': max_attempts}

        self._show_deploy_wait_msg(id, wait_min)
        waiter.wait(deploymentId=id, WaiterConfig=config)

    def _show_deploy_wait_msg(self, id, wait_min):
        sys.stdout.write(
            self.MSG_WAITING.format(deployment_id=id, wait=wait_min)
        )
        sys.stdout.flush()


class CodeDeployValidator:
    def __init__(self, cd_client, resources):
        self._client = cd_client
        self._resource_names = resources

    def describe_cd_resources(self):
        try:
            self.app_details = self._client.get_application(
                applicationName=self._resource_names['app_name']
            )
        except ClientError as e:
            raise exceptions.ServiceClientError(
                action='describe Code Deploy application', error=e
            )

        try:
            dgp = self._resource_names['deployment_group_name']
            app = self._resource_names['app_name']
            self.deployment_group_details = self._client.get_deployment_group(
                applicationName=app, deploymentGroupName=dgp
            )
        except ClientError as e:
            raise exceptions.ServiceClientError(
                action='describe Code Deploy deployment group', error=e
            )

    def get_deployment_wait_time(self):
        if (
            not hasattr(self, 'deployment_group_details')
            or self.deployment_group_details is None
        ):
            return None
        else:
            dgp_info = self.deployment_group_details['deploymentGroupInfo']
            blue_green_info = dgp_info['blueGreenDeploymentConfiguration']

            deploy_ready_wait_min = blue_green_info['deploymentReadyOption'][
                'waitTimeInMinutes'
            ]

            terminate_key = 'terminateBlueInstancesOnDeploymentSuccess'
            termination_wait_min = blue_green_info[terminate_key][
                'terminationWaitTimeInMinutes'
            ]

            configured_wait = deploy_ready_wait_min + termination_wait_min

            return configured_wait + TIMEOUT_BUFFER_MIN

    def validate_all(self):
        self.validate_application()
        self.validate_deployment_group()

    def validate_application(self):
        app_name = self._resource_names['app_name']
        if self.app_details['application']['computePlatform'] != 'ECS':
            raise exceptions.InvalidPlatformError(
                resource='Application', name=app_name
            )

    def validate_deployment_group(self):
        dgp = self._resource_names['deployment_group_name']
        service = self._resource_names['service']
        service_arn = self._resource_names['service_arn']
        cluster = self._resource_names['cluster']
        cluster_arn = self._resource_names['cluster_arn']

        grp_info = self.deployment_group_details['deploymentGroupInfo']
        compute_platform = grp_info['computePlatform']

        if compute_platform != 'ECS':
            raise exceptions.InvalidPlatformError(
                resource='Deployment Group', name=dgp
            )

        target_services = self.deployment_group_details['deploymentGroupInfo'][
            'ecsServices'
        ]

        # either ECS resource names or ARNs can be stored, so check both
        for target in target_services:
            target_serv = target['serviceName']
            if target_serv != service and target_serv != service_arn:
                raise exceptions.InvalidProperyError(
                    dg_name=dgp, resource='service', resource_name=service
                )

            target_cluster = target['clusterName']
            if target_cluster != cluster and target_cluster != cluster_arn:
                raise exceptions.InvalidProperyError(
                    dg_name=dgp, resource='cluster', resource_name=cluster
                )


class ECSClient:
    def __init__(self, session, parsed_args, parsed_globals, user_agent_extra):
        self._args = parsed_args
        self._custom_config = config.Config(user_agent_extra=user_agent_extra)
        self._client = session.create_client(
            'ecs',
            region_name=parsed_globals.region,
            endpoint_url=parsed_globals.endpoint_url,
            verify=parsed_globals.verify_ssl,
            config=self._custom_config,
        )

    def get_service_details(self):
        cluster = self._args.cluster

        if cluster is None or '':
            cluster = 'default'

        try:
            service_response = self._client.describe_services(
                cluster=cluster, services=[self._args.service]
            )
        except ClientError as e:
            raise exceptions.ServiceClientError(
                action='describe ECS service', error=e
            )

        if len(service_response['services']) == 0:
            raise exceptions.InvalidServiceError(
                service=self._args.service, cluster=cluster
            )

        service_details = service_response['services'][0]
        cluster_name = filehelpers.get_cluster_name_from_arn(
            service_details['clusterArn']
        )

        return {
            'service_arn': service_details['serviceArn'],
            'service_name': service_details['serviceName'],
            'cluster_arn': service_details['clusterArn'],
            'cluster_name': cluster_name,
        }

    def register_task_definition(self, kwargs):
        try:
            response = self._client.register_task_definition(**kwargs)
        except ClientError as e:
            raise exceptions.ServiceClientError(
                action='register ECS task definition', error=e
            )

        return response
