import time
import structlog
from httpx import Response, HTTPError
from .client import Plugin

log = structlog.get_logger("careful")


def retry_default_rule(response: Response) -> bool:
    """default behavior is to retry 400s and 500s but not 404s"""
    return response.status_code >= 400 and response.status_code != 404


def retry_only_500s(response: Response) -> bool:
    """retry all status codes that are 500 or above"""
    return response.status_code >= 500


def retry_all_400s_500s(response: Response) -> bool:
    """retry all status codes that are 400 or above"""
    return response.status_code >= 400


class RetryPlugin(Plugin):
    def __init__(
        self,
        attempts: int = 1,
        wait_seconds: float = 10,
        should_retry=retry_default_rule,
    ):
        self._attempts = max(0, attempts)
        self._wait_seconds = wait_seconds
        self._should_retry = should_retry

    def wrap(self, next_handler, client):
        def retry_handler(method: str, url: str, **kwargs) -> Response:
            tries = 0
            exception_raised = None

            while tries <= self._attempts:
                exception_raised = None

                try:
                    tries += 1
                    resp = next_handler(method, url, **kwargs)

                    # break from loop on an accepted response
                    if not self._should_retry(resp):
                        break

                except HTTPError as e:
                    exception_raised = e

                    if exception_response := getattr(e, "response", None):
                        if not self._should_retry(exception_response):
                            break

                # if we're going to retry, sleep first
                if tries <= self._attempts:
                    # twice as long each time
                    wait = self._wait_seconds * (2 ** (tries - 1))
                    if exception_raised:
                        log.info(
                            "retry due to exception",
                            exception=exception_raised,
                            sleep=wait,
                            tries=tries,
                        )
                    else:
                        log.info(
                            "retry due to response rejection",
                            response=resp,
                            sleep=wait,
                            tries=tries,
                        )
                    time.sleep(wait)

            # out of the loop, either an exception was raised or we had a success
            if exception_raised:
                raise exception_raised
            return resp

        return retry_handler
