from collections.abc import Mapping, Sequence
from typing import Any

import aws_cdk as cdk
import constructs
from aibs_informatics_core.env import EnvBase, EnvType, ResourceNameBaseEnum
from aws_cdk import aws_cloudwatch as cloudwatch
from aws_cdk import aws_lambda as lambda_
from aws_cdk import aws_sns as sns
from aws_cdk import aws_sqs as sqs
from aws_cdk.aws_lambda_event_sources import SqsEventSource
from aws_cdk.aws_sns_subscriptions import SqsSubscription


class ExternalSnsTrigger(constructs.Construct):
    """This intended to be a generic CDK construct that defines resources necessary to
    implement a trigger system that listens for events from an external
    (some other AWS account) SNS Topic and fires off a lambda in the
    account where this stack is deployed.

    ``` Diagram:
                                ┌---- ExternalSNSTrigger ----┐
                                |                            |
        (other AWS account)     |           (SQS)            | (optional provided Lambda)
        external_sns_topic -----> external_sns_event_queue -----> triggered_lambda_fn
                                |             |              |
                                |             v              |
                                |         (SQS DLQ)          |
                                |   external_sns_event_dlq   |
                                |             |              |
                                |             v              |
                                |        (Cloudwatch)        |
                                | triggered_lambda_dlq_alarm |
                                |                            |
                                └----------------------------┘

    ```

    The main intended use case is to provide a simple template for setting up automation
    based on PTS SNS Topic notifications in a separate AWS account (e.g. in an informatics processing pipeline)

    Example usage:
    self.external_sns_trigger_construct = ExternalSnsTrigger(
        scope=self,
        id=self.env_base.get_construct_id("merscope-imaging-process-sns-trigger"),
        env_base=self.env_base,
        triggered_lambda_fn=self.pts_listener_fn,
        external_sns_event_name="merscope-imaging-process",
        external_sns_topic_arn=f"arn:aws:sns:us-west-2:{account_id}:{sns_topic_name}",
    )

    NOTE: The `triggered_lambda_fn` is optional and if you have an alternative arrangement
          for triggering off of SQS messages (e.g. Airflow SQS Sensor) you can provide None as the
          `triggered_lambda_fn` when instantiating the ExternalSnsTrigger construct.
    """  # noqa: E501

    @property
    def queue_name(self) -> str:
        if self._external_sns_event_queue_name is not None:
            return self._env_base.prefixed(self._external_sns_event_queue_name)
        return self._env_base.prefixed(self._external_sns_event_name, "sns-event-queue")

    @property
    def dlq_name(self) -> str:
        if self._external_sns_event_dlq_name is not None:
            return self._env_base.prefixed(self._external_sns_event_dlq_name)
        return self._env_base.prefixed(self._external_sns_event_name, "sns-event-dlq")

    def __init__(
        self,
        scope: constructs.Construct,
        id: str,
        env_base: EnvBase,
        triggered_lambda_fn: lambda_.Function | None,
        external_sns_event_name: str | ResourceNameBaseEnum,
        external_sns_topic_arn: str,
        external_sns_event_queue_filters: Sequence[Mapping[str, Any]] | None = None,
        external_sns_event_queue_name: str | None = None,
        external_sns_event_dlq_name: str | None = None,
        external_sns_event_queue_retention_period: cdk.Duration | None = cdk.Duration.days(7),
        sqs_event_source_enabled: bool | None = None,
        **kwargs,
    ) -> None:
        super().__init__(scope=scope, id=id)

        self._env_base = env_base
        self._external_sns_event_name = external_sns_event_name
        self._external_sns_event_queue_name = external_sns_event_queue_name
        self._external_sns_event_dlq_name = external_sns_event_dlq_name

        if sqs_event_source_enabled is None:
            sqs_event_source_enabled = env_base.env_type is EnvType.PROD

        self.external_sns_event_dlq = sqs.Queue(
            scope=self,
            id=env_base.get_construct_id(external_sns_event_name, "sns-event-dlq"),
            queue_name=self.dlq_name,
            retention_period=cdk.Duration.days(14),
        )

        self.external_sns_event_queue = sqs.Queue(
            scope=self,
            id=env_base.get_construct_id(external_sns_event_name, "sns-event-queue"),
            queue_name=self.queue_name,
            retention_period=external_sns_event_queue_retention_period,
            dead_letter_queue=sqs.DeadLetterQueue(
                max_receive_count=2,
                queue=self.external_sns_event_dlq,
            ),
            # visibility_timeout must be longer than the `timeout` of the
            # `triggered_lambda_fn` or deploy will fail
            visibility_timeout=cdk.Duration.seconds(330),
        )

        # The *owning* AWS account of the SNS topic needs to give Subscribe permissions to the
        # AWS account where this stack will be deployed
        # See: https://docs.aws.amazon.com/sns/latest/dg/sns-send-message-to-sqs-cross-account.html
        self.external_sns_topic = sns.Topic.from_topic_arn(
            scope=self,
            id=env_base.get_construct_id(external_sns_event_name, "external-sns-topic"),
            topic_arn=external_sns_topic_arn,
        )

        # Useful reference:
        # https://docs.aws.amazon.com/cdk/api/v2/python/aws_cdk.aws_sns/Subscription.html#aws_cdk.aws_sns.Subscription
        self.external_sns_topic.add_subscription(
            SqsSubscription(
                queue=self.external_sns_event_queue,
                raw_message_delivery=True,
            )
        )

        if triggered_lambda_fn is not None:
            triggered_lambda_fn.add_event_source(
                source=SqsEventSource(
                    queue=self.external_sns_event_queue,
                    report_batch_item_failures=True,
                    filters=external_sns_event_queue_filters,
                    enabled=sqs_event_source_enabled,
                )
            )
            self.external_sns_event_queue.grant_consume_messages(triggered_lambda_fn)

        # Alarm that fires if external_sns_event_queue fails delivery or if lambda fails to process
        # Further actions can be configured by accessing ExternalSnsTrigger.triggered_lambda_dlq_alarm resource # noqa: E501
        # TODO: Revisit making alarms even more configurable in the future
        self.triggered_lambda_dlq_alarm = cloudwatch.Alarm(
            scope=self,
            id=env_base.get_construct_id(external_sns_event_name, "sns-event-dlq-alarm"),
            alarm_description=(
                f"Alarm if more than 1 message in {self.dlq_name} in 10 minute period"
            ),
            metric=self.external_sns_event_dlq.metric_approximate_number_of_messages_visible(
                statistic=cloudwatch.Stats.MAXIMUM,
                period=cdk.Duration.minutes(10),
            ),
            evaluation_periods=1,
            threshold=1,
            datapoints_to_alarm=1,
            comparison_operator=cloudwatch.ComparisonOperator.GREATER_THAN_OR_EQUAL_TO_THRESHOLD,
        )
