package nodemanager

import (
	"context"
	"fmt"
	"time"

	"github.com/pingcap/errors"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
)

const (
	poolManagerRetryLimit   = 10
	poolManagerBackoffLimit = 10 * time.Second
)

type retryablePoolManager struct {
	rm.PoolManager
}

func withRetry[Result any](ctx context.Context, operation func() (Result, error)) (Result, error) {
	var lastErr error
	var zero Result
	backoff := 1 * time.Second
	for i := 0; i < poolManagerRetryLimit; i++ {
		result, err := operation()
		if err == nil {
			return result, nil
		}
		lastErr = err
		select {
		case <-ctx.Done():
			return zero, ctx.Err()
		default:
			time.Sleep(backoff)
		}
		backoff *= 2
		if backoff > poolManagerBackoffLimit {
			backoff = poolManagerBackoffLimit
		}
	}

	return zero, fmt.Errorf("failed after %d retries: %w", poolManagerRetryLimit, lastErr)
}

func (m retryablePoolManager) ActivateWorker(
	ctx context.Context,
	pool rm.WorkerPool,
	worker rm.Worker,
	parameter tidb.ActivateRequest,
) error {
	_, err := withRetry(ctx, func() (bool, error) {
		return true, m.PoolManager.ActivateWorker(ctx, pool, worker, parameter)
	})
	return err
}

func (m retryablePoolManager) RecycleWorker(
	ctx context.Context,
	pool rm.WorkerPool,
	worker rm.Worker,
	keyspaceName string,
) error {
	_, err := withRetry(ctx, func() (bool, error) {
		return true, m.PoolManager.RecycleWorker(ctx, pool, worker, keyspaceName)
	})
	return err
}

func (m retryablePoolManager) GetPools() (map[rm.WorkerPool]*rm.PoolConfig, error) {
	return withRetry(context.Background(), func() (map[rm.WorkerPool]*rm.PoolConfig, error) {
		pools, _ := m.PoolManager.GetPools()
		if len(pools) == 0 {
			return nil, errors.New("no pools found")
		}
		return pools, nil
	})
}

func (m retryablePoolManager) GetStandbyWorkers(pool rm.WorkerPool) ([]rm.Worker, error) {
	return withRetry(context.Background(), func() ([]rm.Worker, error) {
		return m.PoolManager.GetStandbyWorkers(pool)
	})
}