package round_trippers

import (
	"fmt"
	"net/http"
	"time"

	"github.com/rs/zerolog/log"
)

type RetryRoundTripper struct {
	rt              http.RoundTripper
	maxRetries      int
	initialDelaySec int
}

func NewRetryRoundTripper(rt http.RoundTripper, maxRetries int, initialDelaySec int) http.RoundTripper {
	return &RetryRoundTripper{
		rt:              rt,
		maxRetries:      maxRetries,
		initialDelaySec: initialDelaySec,
	}
}

func (rrt *RetryRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
	var resp *http.Response
	var err error

	for attempt := 0; attempt <= rrt.maxRetries; attempt++ {
		if attempt > 0 {
			// exponential retry time (e.g.: [1s, 2s, 4s])
			delay := time.Duration(rrt.initialDelaySec<<uint(attempt-1)) * time.Second

			log.Debug().Msgf("retrying, delay=%d", delay)

			timer := time.NewTimer(delay)

			select {
			case <-timer.C:
			case <-r.Context().Done():
				timer.Stop()

				return nil, r.Context().Err()
			}
		}

		resp, err = rrt.rt.RoundTrip(r)
		if err != nil {
			return nil, err
		}

		switch resp.StatusCode {
		case http.StatusTooManyRequests, 420:
			resp.Body.Close()

			if attempt == rrt.maxRetries {
				return nil, fmt.Errorf("request returned %d after %d retries", resp.StatusCode, rrt.maxRetries)
			}

			continue
		default:
			return resp, nil
		}
	}

	return resp, err
}
