# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License 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.

from __future__ import annotations

import asyncio
import logging
import time
from typing import TYPE_CHECKING, Any

import jmespath
from botocore.exceptions import NoCredentialsError, WaiterError

from airflow.providers.common.compat.sdk import AirflowException

if TYPE_CHECKING:
    from botocore.waiter import Waiter

# Standard throttling and transient error codes to retry on
# https://docs.aws.amazon.com/general/latest/gr/api-retries.html
# and https://github.com/boto/botocore/blob/develop/botocore/retryhandler.py
RETRIABLE_ERROR_CODES = {
    "ThrottlingException",
    "Throttling",
    "RequestLimitExceeded",
    "ProvisionedThroughputExceededException",
    "LimitExceededException",
    "RequestThrottled",
    "RequestThrottledException",
    "TooManyRequestsException",
    "ServerException",
    "InternalServerError",
    "InternalFailure",
    "ServiceUnavailable",
    "BadGateway",
    "GatewayTimeout",
    "RequestTimeout",
    "RequestTimeoutException",
}


def wait(
    waiter: Waiter,
    waiter_delay: int,
    waiter_max_attempts: int,
    args: dict[str, Any],
    failure_message: str,
    status_message: str,
    status_args: list[str],
) -> None:
    """
    Use a boto waiter to poll an AWS service for the specified state.

    Although this function uses boto waiters to poll the state of the
    service, it logs the response of the service after every attempt,
    which is not currently supported by boto waiters.

    :param waiter: The boto waiter to use.
    :param waiter_delay: The amount of time in seconds to wait between attempts.
    :param waiter_max_attempts: The maximum number of attempts to be made.
    :param args: The arguments to pass to the waiter.
    :param failure_message: The message to log if a failure state is reached.
    :param status_message: The message logged when printing the status of the service.
    :param status_args: A list containing the JMESPath queries to retrieve status information from
        the waiter response.
        e.g.
        response = {"Cluster": {"state": "CREATING"}}
        status_args = ["Cluster.state"]

        response = {
        "Clusters": [{"state": "CREATING", "details": "User initiated."},]
        }
        status_args = ["Clusters[0].state", "Clusters[0].details"]
    """
    log = logging.getLogger(__name__)
    first_attempt = True
    attempt = 0
    while attempt < waiter_max_attempts:
        if not first_attempt:
            time.sleep(waiter_delay)
        first_attempt = False
        try:
            waiter.wait(**args, WaiterConfig={"MaxAttempts": 1})

        except NoCredentialsError as error:
            log.info(str(error))

        except WaiterError as error:
            error_reason = str(error)
            last_response = error.last_response

            if "terminal failure" in error_reason:
                log.error("%s: %s", failure_message, _LazyStatusFormatter(status_args, last_response))
                raise AirflowException(f"{failure_message}: {error}")

            if (
                "An error occurred" in error_reason
                and isinstance(last_response.get("Error"), dict)
                and "Code" in last_response.get("Error")
            ):
                error_code = last_response["Error"]["Code"]
                if error_code not in RETRIABLE_ERROR_CODES:
                    raise AirflowException(f"{failure_message}: {error}")

                log.info(
                    "Waiter encountered retriable error: %s. Retrying (attempt %d/%d)...",
                    error_code,
                    attempt + 1,
                    waiter_max_attempts,
                )
                # Don't increment attempt counter for retriable errors; continue looping
                continue

            log.info("%s: %s", status_message, _LazyStatusFormatter(status_args, last_response))
        else:
            break
        attempt += 1
    else:
        raise AirflowException("Waiter error: max attempts reached")


async def async_wait(
    waiter: Waiter,
    waiter_delay: int,
    waiter_max_attempts: int,
    args: dict[str, Any],
    failure_message: str,
    status_message: str,
    status_args: list[str],
) -> None:
    """
    Use an async boto waiter to poll an AWS service for the specified state.

    Although this function uses boto waiters to poll the state of the
    service, it logs the response of the service after every attempt,
    which is not currently supported by boto waiters.

    :param waiter: The boto waiter to use.
    :param waiter_delay: The amount of time in seconds to wait between attempts.
    :param waiter_max_attempts: The maximum number of attempts to be made.
    :param args: The arguments to pass to the waiter.
    :param failure_message: The message to log if a failure state is reached.
    :param status_message: The message logged when printing the status of the service.
    :param status_args: A list containing the JMESPath queries to retrieve status information from
        the waiter response.
        e.g.
        response = {"Cluster": {"state": "CREATING"}}
        status_args = ["Cluster.state"]

        response = {
        "Clusters": [{"state": "CREATING", "details": "User initiated."},]
        }
        status_args = ["Clusters[0].state", "Clusters[0].details"]
    """
    log = logging.getLogger(__name__)
    first_attempt = True
    attempt = 0
    while attempt < waiter_max_attempts:
        if not first_attempt:
            await asyncio.sleep(waiter_delay)
        first_attempt = False
        try:
            await waiter.wait(**args, WaiterConfig={"MaxAttempts": 1})

        except NoCredentialsError as error:
            log.info(str(error))

        except WaiterError as error:
            error_reason = str(error)
            last_response = error.last_response

            if "terminal failure" in error_reason:
                raise AirflowException(
                    f"{failure_message}: {_LazyStatusFormatter(status_args, last_response)}\n{error}"
                )

            if (
                "An error occurred" in error_reason
                and isinstance(last_response.get("Error"), dict)
                and "Code" in last_response.get("Error")
            ):
                error_code = last_response["Error"]["Code"]
                if error_code not in RETRIABLE_ERROR_CODES:
                    raise AirflowException(f"{failure_message}\n{last_response}\n{error}")

                log.info(
                    "Waiter encountered retriable error: %s. Retrying (attempt %d/%d)...",
                    error_code,
                    attempt + 1,
                    waiter_max_attempts,
                )
                # Don't increment attempt counter for retriable errors; continue looping
                continue

            log.info("%s: %s", status_message, _LazyStatusFormatter(status_args, last_response))
        else:
            break
        attempt += 1
    else:
        raise AirflowException("Waiter error: max attempts reached")


class _LazyStatusFormatter:
    """
    Contains the info necessary to extract the status from a response; only computes the value when necessary.

    Used to avoid computations if the logs are disabled at the given level.
    """

    def __init__(self, jmespath_queries: list[str], response: dict[str, Any]):
        self.jmespath_queries = jmespath_queries
        self.response = response

    def __str__(self):
        """Loop through the args list and generate a string containing values from the waiter response."""
        values = []
        for query in self.jmespath_queries:
            value = jmespath.search(query, self.response)
            if value is not None and value != "":
                values.append(str(value))

        return " - ".join(values)
