#!/usr/bin/env python

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, argparse, inspect, json, hashlib, time, copy, shutil

import attr
import botocore, boto3.session
import chalice, chalice.app, chalice.awsclient, chalice.deploy.packager, chalice.deploy.deployer, chalice.deploy.models
from chalice.cli.factory import create_botocore_session, CLIFactory
from chalice.deploy.packager import LambdaDeploymentPackager
from chalice.deploy.deployer import create_default_deployer
from chalice.utils import UI
from chalice.compat import urlparse
from chalice.constants import DEFAULT_STAGE_NAME

import domovoi

try:
    import pkg_resources
    __version__ = pkg_resources.get_distribution("domovoi").version
except Exception:
    __version__ = "0.0.0"

class DomovoiDeploymentPackager(LambdaDeploymentPackager):
    _CHALICE_LIB_DIR = "domovoilib"
    def _add_app_files(self, zip_fileobj, project_dir):
        domovoi_router = inspect.getfile(domovoi.app)
        if domovoi_router.endswith(".pyc"):
            domovoi_router = domovoi_router[:-1]
        zip_fileobj.write(domovoi_router, "domovoi/app.py")

        domovoi_init = inspect.getfile(domovoi)
        if domovoi_init.endswith(".pyc"):
            domovoi_init = domovoi_init[:-1]
        zip_fileobj.write(domovoi_init, "domovoi/__init__.py")

        chalice_router = inspect.getfile(chalice.app)
        if chalice_router.endswith(".pyc"):
            chalice_router = chalice_router[:-1]
        zip_fileobj.write(chalice_router, "chalice/app.py")

        chalice_init = inspect.getfile(chalice)
        if chalice_init.endswith(".pyc"):
            chalice_init = chalice_init[:-1]
        zip_fileobj.write(chalice_init, "chalice/__init__.py")

        zip_fileobj.write(os.path.join(project_dir, "app.py"), "app.py")
        self._add_chalice_lib_if_needed(project_dir, zip_fileobj)

    def _needs_latest_version(self, filename):
        return filename == 'app.py' or filename.startswith(('domovoilib/', 'domovoi/'))

    def create_deployment_package(self, project_dir, python_version, package_filename=None):
        deployment_package_filename = self.deployment_package_filename(project_dir, python_version)
        if os.path.exists(deployment_package_filename):
            self.inject_latest_app(deployment_package_filename, project_dir)
            return deployment_package_filename
        else:
            return LambdaDeploymentPackager.create_deployment_package(self, project_dir, python_version,
                                                                      package_filename=package_filename)

@attr.attrs
class ManagedIAMRole(chalice.deploy.models.ManagedIAMRole):
    def __attrs_post_init__(self):
        self.role_name = self.role_name.rpartition("-")[0]

@attr.attrs
class LambdaFunction(chalice.deploy.models.LambdaFunction):
    def __attrs_post_init__(self):
        self.function_name = self.function_name.rpartition("-")[0]


chalice.deploy.packager.LambdaDeploymentPackager = DomovoiDeploymentPackager
chalice.deploy.deployer.LambdaDeploymentPackager = DomovoiDeploymentPackager
chalice.deploy.models.ManagedIAMRole = ManagedIAMRole
chalice.deploy.models.LambdaFunction = LambdaFunction

def add_filter_config(event_config, event_handler):
    cfg = dict(event_config)
    for fltr in "prefix", "suffix":
        if event_handler.get(fltr):
            cfg.setdefault("Filter", dict(Key=dict(FilterRules=[])))
            cfg["Filter"]["Key"]["FilterRules"].append(dict(Name=fltr, Value=event_handler[fltr]))
    return cfg

def put_event_source_mapping(event_source_arn, function_name, source_data, dry_run=False):
    event_source_mapping_args = dict(EventSourceArn=event_source_arn,
                                     FunctionName=function_name,
                                     Enabled=True)
    if "dynamodb" in event_source_arn:
        event_source_mapping_args.update(StartingPosition="TRIM_HORIZON")
    if source_data["batch_size"] is not None:
        event_source_mapping_args.update(BatchSize=source_data["batch_size"])
    esm = None
    try:
        if not dry_run:
            esm = awslambda.create_event_source_mapping(**event_source_mapping_args)
    except awslambda.exceptions.ResourceConflictException as e:
        assert "already exists" in str(e) and str(e).split()[-2] == "UUID"
        if source_data["batch_size"] is not None:
            esm_uuid = str(e).split()[-1]
            esm = awslambda.get_event_source_mapping(UUID=esm_uuid)
            if source_data["batch_size"] != esm["BatchSize"]:
                esm = awslambda.update_event_source_mapping(UUID=esm_uuid, BatchSize=source_data["batch_size"])
    return esm


parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("--stage", default=DEFAULT_STAGE_NAME)
parser.add_argument("--profile")
parser.add_argument("--debug", action="store_true")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--version", action="version", version="domovoi {}".format(__version__))
parser.add_argument("action", choices={"deploy", "new-project"})
parser.add_argument("project_dir", nargs="?", default=os.getcwd())
args = parser.parse_args()

example_app = """# This is an example entry point for an app built with Domovoi, an AWS Lambda event handler manager.
# See https://github.com/kislyuk/domovoi for Domovoi's documentation.

import json, boto3, domovoi

app = domovoi.Domovoi()

@app.scheduled_function("rate(1 minute)")
def test(event, context):
    pass
"""

if args.action == "new-project":
    from chalice.cli import click, create_new_project_skeleton
    from chalice.policy import PolicyBuilder
    if args.project_dir == os.getcwd():
        args.project_dir = click.prompt("New Domovoi project name")
    assert not os.path.isdir(args.project_dir)
    create_new_project_skeleton(args.project_dir, profile=None)
    with open(os.path.join(args.project_dir, "app.py"), "w") as fh:
        fh.write(example_app)
    default_iam_policy_filename = os.path.join(os.path.dirname(domovoi.__file__), "default_iam_policy.json")
    stage_iam_policy_filename = os.path.join(args.project_dir, ".chalice", "policy-{}.json".format(DEFAULT_STAGE_NAME))
    shutil.copy(default_iam_policy_filename, stage_iam_policy_filename)
    parser.exit(status=0, message="New Domovoi project created in {}\n".format(args.project_dir))

botocore_session = create_botocore_session(profile=args.profile, debug=args.debug)
botocore_session.user_agent_extra = "domovoi/%s" % __version__
boto3_session = boto3.session.Session(botocore_session=botocore_session)
events = boto3_session.client("events")
sns = boto3_session.resource("sns")
sqs = boto3_session.resource("sqs")
awslambda = boto3_session.client("lambda")
s3 = boto3_session.resource("s3")
sts = boto3_session.client("sts")
sfn = boto3.client("stepfunctions")
iam = boto3_session.resource("iam")
logs = boto3_session.client("logs")
dynamodb = boto3_session.resource("dynamodb")
ec2 = boto3_session.resource("ec2")
elbv2 = boto3_session.client("elbv2")
acm = boto3_session.client("acm")
route53 = boto3_session.client("route53")

trust_statement = copy.deepcopy(chalice.constants.LAMBDA_TRUST_POLICY["Statement"][0])
trust_statement["Principal"] = {"Service": urlparse(sfn.meta.endpoint_url).netloc}
chalice.constants.LAMBDA_TRUST_POLICY["Statement"].append(trust_statement)

config = CLIFactory(args.project_dir).create_config_obj(chalice_stage_name=args.stage, autogen_policy=False)
config._user_provided_params.setdefault("tags", {})
config._user_provided_params["tags"]["domovoi"] = "version={}".format(__version__)
deployer = create_default_deployer(session=botocore_session,
                                   config=config,
                                   ui=UI(confirm=lambda *args, **kwargs: True))
function_name = '%s-%s' % (config.app_name, args.stage)

if args.dry_run:
    deployed_values, lambda_arn = dict(resources=[]), "arn:aws:lambda:::"
else:
    deployed_values = deployer.deploy(config, chalice_stage_name=args.stage)
    for resource in deployed_values["resources"]:
        if resource["resource_type"] == "lambda_function":
            lambda_arn = resource["lambda_arn"]

    client = chalice.awsclient.TypedAWSClient(botocore_session)._client("lambda")
    fn_updates = dict(Description=config._chain_lookup("description") or "Domovoi event handler")
    dlq_arn = config._chain_lookup('dead_letter_queue_target_arn')
    if dlq_arn:
        print("Setting DLQ {} for {}".format(dlq_arn, function_name))
        fn_updates.update(DeadLetterConfig=dict(TargetArn=dlq_arn))
    client.update_function_configuration(FunctionName=function_name, **fn_updates)
    reserved_concurrent_executions = config._chain_lookup('reserved_concurrent_executions')
    if reserved_concurrent_executions:
        print("Setting concurrency reservation {} for {}".format(reserved_concurrent_executions, function_name))
        client.put_function_concurrency(FunctionName=function_name,
                                        ReservedConcurrentExecutions=reserved_concurrent_executions)
    alb_acm_cert_dns_name = config._chain_lookup('alb_acm_cert_dns_name')

domovoi_app = config.chalice_app

# TODO: consider narrowing trust policy
for service in ("apigateway", "events", "sns", "sqs", "s3", "logs", "elasticloadbalancing"):
    service_uri = service + ".amazonaws.com"
    policy = dict(FunctionName=lambda_arn, Principal=service_uri, Action="lambda:InvokeFunction")
    policy_id = "domovoi-{}".format(hashlib.md5(json.dumps(policy).encode()).hexdigest()[:8])
    print("Granting {} access to invoke Lambda function {}".format(service, lambda_arn))
    if not args.dry_run:
        try:
            awslambda.add_permission(StatementId=policy_id, **policy)
        except awslambda.exceptions.ResourceConflictException:
            print("Found existing permission grant statement {}, skipping".format(policy_id))

def find_acm_cert(dns_name):
    for page in acm.get_paginator("list_certificates").paginate(CertificateStatuses=["ISSUED"]):
        for cert in page["CertificateSummaryList"]:
            if cert["DomainName"] == dns_name:
                return cert
    raise Exception("Unable to find ACM certificate for {}".format(dns_name))

def ensure_ingress_rule(security_group, **kwargs):
    cidr_ip = kwargs.pop("CidrIp")
    for rule in security_group.ip_permissions:
        ip_range_matches = any(cidr_ip == ip_range["CidrIp"] for ip_range in rule["IpRanges"])
        opts_match = all(rule.get(arg) == kwargs[arg] for arg in kwargs)
        if ip_range_matches and opts_match:
            break
    else:
        security_group.authorize_ingress(CidrIp=cidr_ip, **kwargs)

def find_route53_zone_for_name(dns_name):
    best_zone = None
    for page in route53.get_paginator("list_hosted_zones").paginate():
        for zone in page["HostedZones"]:
            if dns_name.endswith(zone["Name"].rstrip(".")):
                if best_zone is None or len(best_zone["Name"]) < len(zone["Name"]):
                    best_zone = zone
    return best_zone

def update_alias_dns_record(source_dns_name, target_dns_name, hosted_zone_id):
    zone = find_route53_zone_for_name(source_dns_name)
    zone_update = {
        "Action": "UPSERT",
        "ResourceRecordSet": {
            "Name": source_dns_name,
            "Type": "A",
            "AliasTarget": {
                'HostedZoneId': hosted_zone_id,
                'DNSName': target_dns_name,
                'EvaluateTargetHealth': False
            }
        }
    }
    route53.change_resource_record_sets(HostedZoneId=zone["Id"], ChangeBatch=dict(Changes=[zone_update]))

def register_deployed_resource(resource_type, **kwargs):
    deployed_values["resources"].append(dict(kwargs, resource_type=resource_type))
    deployer._recorder.record_results(deployed_values, args.stage, args.project_dir)


if domovoi_app.alb_targets:
    if not args.dry_run:
        if not alb_acm_cert_dns_name:
            raise Exception('Please set the "alb_acm_cert_dns_name" config key in your .chalice/config.json '
                            'to the DNS name of a validated ACM certificate in your account')
        for vpc in ec2.vpcs.filter(Filters=[dict(Name="isDefault", Values=["true"])]):
            break
        else:
            raise Exception("A default VPC is required")
        security_group_name = "domovoi-{}".format(function_name)
        try:
            security_group = vpc.create_security_group(GroupName=security_group_name,
                                                       Description="Automatically managed by Domovoi for Lambda ALB")
        except botocore.exceptions.ClientError as e:
            if "InvalidGroup.Duplicate" not in str(e):
                raise
            security_group = list(vpc.security_groups.filter(GroupNames=[security_group_name]))[0]

        ensure_ingress_rule(security_group, IpProtocol="tcp", FromPort=443, ToPort=443, CidrIp="0.0.0.0/0")

        for prefix, handler in domovoi_app.alb_targets.items():
            res = elbv2.create_load_balancer(Name=function_name,
                                             Subnets=[subnet.id for subnet in vpc.subnets.all()],
                                             SecurityGroups=[security_group.id])
            alb = res["LoadBalancers"][0]
            print("Using ALB", alb["LoadBalancerArn"])
            res = elbv2.create_target_group(Name=function_name, TargetType="lambda")
            target_group = res["TargetGroups"][0]
            print("Using target group", target_group["TargetGroupArn"])
            cert = find_acm_cert(alb_acm_cert_dns_name)
            print("Using ACM certificate", cert["CertificateArn"])
            default_action = dict(Type="forward", TargetGroupArn=target_group["TargetGroupArn"])
            res = elbv2.create_listener(LoadBalancerArn=alb["LoadBalancerArn"],
                                        Protocol="HTTPS",
                                        Port=443,
                                        Certificates=[dict(CertificateArn=cert["CertificateArn"])],
                                        DefaultActions=[default_action])
            listener = res["Listeners"][0]
            print("Using listener", listener["ListenerArn"])
            res = elbv2.register_targets(TargetGroupArn=target_group["TargetGroupArn"], Targets=[dict(Id=lambda_arn)])
            print("Updating the Route53 ALIAS DNS record for {} to {}".format(alb_acm_cert_dns_name, alb["DNSName"]))
            update_alias_dns_record(alb_acm_cert_dns_name, alb["DNSName"], hosted_zone_id=alb["CanonicalHostedZoneId"])

for task_name, task in domovoi_app.cloudwatch_events_rules.items():
    print("Scheduling", task_name, "to run on schedule", task["schedule_expression"], "pattern", task["event_pattern"])
    rule_args = dict(Name=task_name)
    if task.get("schedule_expression"):
        rule_args["ScheduleExpression"] = task["schedule_expression"]
    if task.get("event_pattern"):
        rule_args["EventPattern"] = json.dumps(task["event_pattern"])
    if not args.dry_run:
        rule_arn = events.put_rule(**rule_args)["RuleArn"]
        lambda_input = '{"task_name": "%s", "event": <event>}' % task_name
        ixform = dict(InputPathsMap=dict(event="$"), InputTemplate=lambda_input)
        events.put_targets(Rule=task_name, Targets=[dict(Id=task_name, Arn=lambda_arn, InputTransformer=ixform)])
        print("Scheduled CloudWatch event", rule_arn)

for sns_topic, event_handler in domovoi_app.sns_subscribers.items():
    print("Subscribing", event_handler, "to SNS topic", sns_topic)
    if not args.dry_run:
        topic = sns.create_topic(Name=sns_topic)
        subscription = topic.subscribe(Protocol="lambda", Endpoint=lambda_arn)
        print("Subscribed to", subscription)

for sqs_queue, source_data in domovoi_app.sqs_subscribers.items():
    print("Subscribing", source_data["func"], "to SQS queue", sqs_queue)
    queue_attributes = dict(domovoi_app.sqs_default_queue_attributes)
    queue_attributes.update(source_data["queue_attributes"] or {})
    if not args.dry_run:
        queue = sqs.create_queue(QueueName=sqs_queue)
        queue_arn = queue.attributes["QueueArn"]
        queue.set_attributes(Attributes=queue_attributes)
        put_event_source_mapping(event_source_arn=queue_arn,
                                 function_name=function_name,
                                 source_data=source_data,
                                 dry_run=args.dry_run)

def ensure_queue(queue_name, sender_arn, event_handler):
    sqs_queue = sqs.create_queue(QueueName=queue_name)
    policy = {"Statement": [{"Action": ["SQS:SendMessage"],
                             "Effect": "Allow",
                             "Resource": sqs_queue.attributes["QueueArn"],
                             "Principal": {"AWS": "*"},
                             "Condition": {"ArnLike": {"aws:SourceArn": sender_arn}}}]}
    policy["Statement"][0]["Sid"] = "domovoi-{}".format(hashlib.md5(json.dumps(policy).encode()).hexdigest()[:8])
    queue_attributes = dict(domovoi_app.sqs_default_queue_attributes)
    queue_attributes.update(event_handler["sqs_queue_attributes"] or {}, Policy=json.dumps(policy))
    sqs_queue.set_attributes(Attributes=queue_attributes)
    return sqs_queue

def ensure_topic(topic_name, s3_bucket):
    sns_topic = sns.create_topic(Name=topic_name)
    policy = {"Statement": [{"Action": ["SNS:Publish"],
                             "Effect": "Allow",
                             "Resource": sns_topic.arn,
                             "Principal": {"Service": ["s3.amazonaws.com"]},
                             "Condition": {"ArnLike": {"aws:SourceArn": "arn:aws:s3:*:*:" + s3_bucket}}}]}
    policy["Statement"][0]["Sid"] = "domovoi-{}".format(hashlib.md5(json.dumps(policy).encode()).hexdigest()[:8])
    sns_topic.set_attributes(AttributeName="Policy", AttributeValue=json.dumps(policy))
    return sns_topic


for s3_bucket, event_handler in domovoi_app.s3_subscribers.items():
    print("Subscribing", event_handler["func"], "to events in S3 bucket", s3_bucket)
    if args.dry_run:
        continue
    if event_handler["use_sqs"] and event_handler["use_sns"]:
        # An SNS-SQS bridge is the only option for subscribing multiple Lambdas to the same S3 event type with SQS
        topic_name = "domovoi-s3-events-{}".format(s3_bucket.replace(".", "_"))
        sns_topic = ensure_topic(topic_name, s3_bucket)
        queue_name = "domovoi-s3-events-{}-{}".format(s3_bucket.replace(".", "_"), function_name)
        sqs_queue = ensure_queue(queue_name, sender_arn="arn:aws:sns:*:*:" + topic_name, event_handler=event_handler)
        queue_arn = sqs_queue.attributes["QueueArn"]
        subscription = sns_topic.subscribe(Protocol="sqs", Endpoint=queue_arn)
        print("Subscribed", queue_arn, "to", subscription)
        esm = put_event_source_mapping(event_source_arn=queue_arn,
                                       function_name=function_name,
                                       source_data=dict(batch_size=event_handler["sqs_batch_size"]))
        print("Created event source mapping", esm["UUID"])
        topic_configuration = dict(TopicArn=sns_topic.arn, Events=event_handler["events"])
        topic_configuration = add_filter_config(topic_configuration, event_handler)
    elif event_handler["use_sqs"]:
        queue_name = "domovoi-s3-events-{}-{}".format(s3_bucket.replace(".", "_"), function_name)
        sqs_queue = ensure_queue(queue_name, sender_arn="arn:aws:s3:*:*:" + s3_bucket, event_handler=event_handler)
        queue_arn = sqs_queue.attributes["QueueArn"]
        esm = put_event_source_mapping(event_source_arn=queue_arn,
                                       function_name=function_name,
                                       source_data=dict(batch_size=event_handler["sqs_batch_size"]))
        print("Created event source mapping", esm["UUID"])
        queue_configuration = dict(QueueArn=queue_arn, Events=event_handler["events"])
        queue_configuration = add_filter_config(queue_configuration, event_handler)
    elif event_handler["use_sns"]:
        topic_name = "domovoi-s3-events-{}".format(s3_bucket.replace(".", "_"))
        sns_topic = ensure_topic(topic_name, s3_bucket)
        subscription = sns_topic.subscribe(Protocol="lambda", Endpoint=lambda_arn)
        print("Subscribed", lambda_arn, "to", subscription)
        topic_configuration = dict(TopicArn=sns_topic.arn, Events=event_handler["events"])
        topic_configuration = add_filter_config(topic_configuration, event_handler)
    else:
        lambda_function_configuration = dict(LambdaFunctionArn=lambda_arn, Events=event_handler["events"])
        lambda_function_configuration = add_filter_config(lambda_function_configuration, event_handler)
    for t in range(8):
        try:
            notification, last_exception = s3.Bucket(s3_bucket).Notification(), None
            new_config = dict(LambdaFunctionConfigurations=notification.lambda_function_configurations or [],
                              QueueConfigurations=notification.queue_configurations or [],
                              TopicConfigurations=notification.topic_configurations or [])
            if event_handler["use_sqs"] and not event_handler["use_sns"]:
                old_cfgs = [cfg for cfg in new_config["QueueConfigurations"] if cfg["QueueArn"] != queue_arn]
                new_config["QueueConfigurations"] = [queue_configuration] + old_cfgs
            elif event_handler["use_sns"]:
                old_cfgs = [cfg for cfg in new_config["TopicConfigurations"] if cfg["TopicArn"] != sns_topic.arn]
                new_config["TopicConfigurations"] = [topic_configuration] + old_cfgs
            else:
                old_cfgs = [cfg for cfg in new_config["LambdaFunctionConfigurations"]
                            if cfg["LambdaFunctionArn"] != lambda_arn]
                new_config["LambdaFunctionConfigurations"] = [lambda_function_configuration] + old_cfgs
            notification.put(NotificationConfiguration=new_config)
            break
        except botocore.exceptions.ClientError as e:
            if "A conflicting conditional operation is currently in progress" not in str(e):
                raise
            last_exception = e
            print("Waiting", int(1.6**t), "seconds for concurrent operation to complete")
            time.sleep(1.6**t)
    else:
        raise last_exception

for cwl_log_group_name, cwl_sub_filter_data in domovoi_app.cwl_sub_filters.items():
    print("Subscribing", cwl_sub_filter_data, "to CloudWatch Logs filter for", cwl_log_group_name)
    if not args.dry_run:
        logs.put_subscription_filter(logGroupName=cwl_log_group_name,
                                     filterName="domovoi-cwl-filter",
                                     filterPattern=cwl_sub_filter_data["filter_pattern"],
                                     destinationArn=lambda_arn)

for table_name, source_data in domovoi_app.dynamodb_event_sources.items():
    print("Subscribing to DynamoDB event stream for table", table_name)
    stream_arn = dynamodb.Table(table_name).latest_stream_arn if not args.dry_run else "arn:aws::::"
    put_event_source_mapping(event_source_arn=stream_arn,
                             function_name=function_name,
                             source_data=source_data,
                             dry_run=args.dry_run)

existing_aliases = []
if not args.dry_run:
    for page in awslambda.get_paginator('list_aliases').paginate(FunctionName=function_name):
        existing_aliases.extend(page["Aliases"])

state_machine = None
for sfn_task_name, sfn_task in domovoi_app.sfn_tasks.items():
    print("Registering step function state machine for", sfn_task_name)
    if state_machine is None:
        state_machine = sfn_task["state_machine_definition"]
    else:
        msg = "Multiple state machine definitions are not supported"
        assert state_machine == sfn_task["state_machine_definition"], msg
    lambda_alias = "domovoi-stepfunctions-task-" + sfn_task_name
    alias_args = dict(FunctionName=function_name,
                      Name=lambda_alias,
                      FunctionVersion="$LATEST",
                      Description="Domovoi Lambda routing label for a Step Functions state machine task")
    all_states = domovoi.Domovoi.get_all_states(state_machine)
    state = all_states[sfn_task["state_name"]]
    if not args.dry_run:
        for alias in existing_aliases:
            if alias["Name"] == lambda_alias and alias["FunctionVersion"] == "$LATEST":
                break
        else:
            try:
                awslambda.create_alias(**alias_args)
            except awslambda.exceptions.ResourceConflictException:
                awslambda.update_alias(**alias_args)
        state["Resource"] = lambda_arn + ":" + lambda_alias

if state_machine and not args.dry_run:
    iam_role_arn = config.iam_role_arn or iam.Role(function_name).arn
    sm_args = dict(name=function_name,
                   definition=json.dumps(state_machine),
                   roleArn=iam_role_arn)
    try:
        sm = sfn.create_state_machine(**sm_args)
        print("Created new state machine", sm["stateMachineArn"])
    except botocore.exceptions.ClientError as e:
        for page in sfn.get_paginator("list_state_machines").paginate():
            for sm in page["stateMachines"]:
                if sm["name"] == function_name:
                    break
        if sm["name"] != function_name:
            raise e
        sm = sfn.describe_state_machine(stateMachineArn=sm["stateMachineArn"])
        sm_args.clear()
        if json.loads(sm["definition"]) != state_machine:
            sm_args["definition"] = json.dumps(state_machine)
        if sm["roleArn"] != iam_role_arn:
            sm_args["roleArn"] = iam_role_arn
        if sm_args:
            print("Updating state machine", sm["stateMachineArn"])
            sfn.update_state_machine(stateMachineArn=sm["stateMachineArn"], **sm_args)
        else:
            print("No changes required to existing state machine", sm["stateMachineArn"])
    register_deployed_resource("state_machine", name=function_name, arn=sm["stateMachineArn"])
    print("State machine:", sm["stateMachineArn"])

if args.dry_run:
    print("Dry run successful")
