import backoff
import boto3
from rich import print

from airflow_tools import config


class MWAAConnector:
    """Connect to MWAA."""

    def __init__(self, retry_interval: int = 30, timeout: int = 3600):
        self._client = boto3.client("mwaa", region_name=config.AWS_REGION)
        self._retry_interval = retry_interval
        self._timeout = timeout

    def get_retry_interval(self):
        return self._retry_interval

    def get_timeout(self):
        return self._timeout

    def update_environment(
        self,
        environment_name: str,
        bucket: str,
        dag_s3_path: str,
        requirements_s3_path: str,
    ):
        @backoff.on_predicate(
            backoff.constant,
            jitter=None,
            interval=self.get_retry_interval,
            max_time=self.get_timeout,
        )
        def _wait_for_environment_availability(environment_name: str):
            response = self._client.get_environment(Name=environment_name)
            status = response["Environment"]["Status"]

            if status not in ["AVAILABLE", "UPDATE_FAILED", "UNAVAILABLE"]:
                print(
                    f"Environment {environment_name} has updates pending. "
                    f"Current status: {status}"
                )
                return False

            return True

        @backoff.on_predicate(
            backoff.constant,
            jitter=None,
            interval=self.get_retry_interval,
            max_time=self.get_timeout,
        )
        def _wait_for_update_to_complete(environment_name: str):
            response = self._client.get_environment(Name=environment_name)
            status = response["Environment"]["LastUpdate"]["Status"]

            if status == "PENDING":
                print(f"Waiting for update to complete. Current status: {status}")
                return False

            if status == "FAILED":
                error_code = response["Environment"]["LastUpdate"]["Error"]["ErrorCode"]
                error_message = response["Environment"]["LastUpdate"]["Error"][
                    "ErrorMessage"
                ]
                raise Exception(
                    f"Deployment failed. Error code: {error_code}, "
                    f"message: {error_message}"
                )

            return True

        available = _wait_for_environment_availability(environment_name)

        if not available:
            raise Exception(
                f"Timed out waiting for environment {environment_name} to be available."
            )

        print(
            f"Updating environment {environment_name} from bucket {bucket}. "
            f"Dags path: {dag_s3_path}. Requirements path: {requirements_s3_path}"
        )
        self._client.update_environment(
            Name=environment_name,
            SourceBucketArn=f"arn:aws:s3:::{bucket}",
            DagS3Path=dag_s3_path,
            RequirementsS3Path=requirements_s3_path,
        )

        result = _wait_for_update_to_complete(environment_name)

        if not result:
            raise Exception("Timed out waiting for deployment to complete.")