package zarinhub

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"time"
)

type RetryTransport struct {
	base       http.RoundTripper
	timeout    time.Duration
	maxRetries int
	retryWait  func(ctx context.Context, delay time.Duration) bool
	logger     *slog.Logger
}

func NewRetryTransport(
	base http.RoundTripper,
	timeout time.Duration,
	maxRetries int,
	loggers ...*slog.Logger,
) *RetryTransport {
	if base == nil {
		base = http.DefaultTransport
	}
	if maxRetries < 0 {
		maxRetries = 0
	}

	var logger *slog.Logger
	if len(loggers) > 0 && loggers[0] != nil {
		logger = loggers[0]
	}

	return &RetryTransport{
		base:       base,
		timeout:    timeout,
		maxRetries: maxRetries,
		retryWait:  waitForRetry,
		logger:     logger,
	}
}

func NewHTTPClient(
	timeout time.Duration,
	maxRetries int,
	loggers ...*slog.Logger,
) *http.Client {
	transport := NewRetryTransport(
		http.DefaultTransport,
		timeout,
		maxRetries,
		loggers...,
	)

	return &http.Client{
		Transport: transport,
	}
}

func (t *RetryTransport) log() *slog.Logger {
	if t.logger != nil {
		return t.logger
	}
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

func (t *RetryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	waitFn := t.retryWait
	if waitFn == nil {
		waitFn = waitForRetry
	}

	if req.Body != nil && req.GetBody == nil {
		bodyBytes, err := io.ReadAll(req.Body)
		if err != nil {
			return nil, err
		}
		_ = req.Body.Close()
		req.GetBody = func() (io.ReadCloser, error) {
			return io.NopCloser(bytes.NewReader(bodyBytes)), nil
		}
	}

	var (
		resp *http.Response
		err  error
	)

	for attempt := 0; attempt <= t.maxRetries; attempt++ {
		if attempt > 0 {
			t.log().InfoContext(
				req.Context(),
				fmt.Sprintf("[RETRY] zarinhub request attempt %d", attempt),
				"is_retry", true,
				"retry_count", attempt,
			)
		} else {
			t.log().InfoContext(
				req.Context(),
				"[INITIAL] zarinhub request",
				"is_retry", false,
				"retry_count", 0,
			)
		}

		var (
			attemptCtx    context.Context
			cancelAttempt context.CancelFunc
		)
		if t.timeout > 0 {
			attemptCtx, cancelAttempt = context.WithTimeout(req.Context(), t.timeout)
		} else {
			attemptCtx = req.Context()
			cancelAttempt = func() {}
		}

		attemptReq := req.Clone(attemptCtx)
		if req.GetBody != nil {
			body, bodyErr := req.GetBody()
			if bodyErr != nil {
				cancelAttempt()
				return nil, bodyErr
			}
			attemptReq.Body = body
		}

		resp, err = t.base.RoundTrip(attemptReq)

		if req.Context().Err() != nil {
			cancelAttempt()
			if resp != nil && resp.Body != nil {
				_ = resp.Body.Close()
			}
			if errors.Is(req.Context().Err(), context.Canceled) {
				return nil, fmt.Errorf("zarinhub request canceled: %w", req.Context().Err())
			}
			return nil, req.Context().Err()
		}

		isErrRetryable := false
		var retryReason error

		if err != nil {
			isErrRetryable = isTransportErrorRetryable(err)
			retryReason = err
		} else if isStatusCodeRetryable(resp.StatusCode) {
			isErrRetryable = true
			retryReason = fmt.Errorf("upstream returned status %d", resp.StatusCode)
		}

		if !isErrRetryable {
			if resp != nil && resp.Body != nil {
				resp.Body = &bodyWithCancel{
					ReadCloser: resp.Body,
					cancel:     cancelAttempt,
				}
			} else {
				cancelAttempt()
			}
			return resp, err
		}

		if attempt == t.maxRetries {
			if attempt > 0 {
				t.log().ErrorContext(
					req.Context(),
					"[RETRY] zarinhub request failed: all retries exhausted",
					"is_retry", true,
					"retry_count", attempt,
					"error", retryReason,
				)
			}
			if resp != nil && resp.Body != nil {
				resp.Body = &bodyWithCancel{
					ReadCloser: resp.Body,
					cancel:     cancelAttempt,
				}
			} else {
				cancelAttempt()
			}
			return resp, err
		}

		if resp != nil && resp.Body != nil {
			_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
			_ = resp.Body.Close()
		}
		cancelAttempt()

		delay := retryDelay(attempt)
		t.log().WarnContext(
			req.Context(),
			fmt.Sprintf("[RETRY] zarinhub request failed, retrying in %v (retry %d)", delay, attempt+1),
			"is_retry", true,
			"next_retry", attempt+1,
			"delay_ms", delay.Milliseconds(),
			"error", retryReason,
		)

		if !waitFn(req.Context(), delay) {
			if req.Context().Err() != nil {
				if errors.Is(req.Context().Err(), context.Canceled) {
					return nil, fmt.Errorf("zarinhub request canceled: %w", req.Context().Err())
				}
				return nil, req.Context().Err()
			}
			return nil, context.Canceled
		}
	}

	return resp, err
}

func isTransportErrorRetryable(err error) bool {
	return !errors.Is(err, context.Canceled)
}

func isStatusCodeRetryable(statusCode int) bool {
	return statusCode == http.StatusRequestTimeout ||
		statusCode == http.StatusTooManyRequests ||
		statusCode >= http.StatusInternalServerError
}

type bodyWithCancel struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (b *bodyWithCancel) Close() error {
	err := b.ReadCloser.Close()
	b.cancel()
	return err
}