package regionalserver

import (
	"sync"
	"time"

	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

var (
	defaultTimeout                   = 99999 * time.Second
	defaultPerRegionResultBufferSize = 100
)

type RegionalResult struct {
	Items []any
	Err   bizerr.BizError
}

// QueryMultiRegional executes the queryFunc for each region concurrently and returns a channel of RegionalResult.
func QueryMultiRegional(
	regions []*entity.Region,
	queryFunc func(*entity.Region, *sync.WaitGroup, chan<- RegionalResult),
) <-chan RegionalResult {
	var wg sync.WaitGroup
	regionalResultCh := make(chan RegionalResult, len(regions)*defaultPerRegionResultBufferSize)
	doneCh := make(chan struct{})

	for _, region := range regions {
		wg.Add(1)
		go func(region *entity.Region) {
			defer wg.Done()
			queryFunc(region, &wg, regionalResultCh)
		}(region)
	}

	go func() {
		defer close(regionalResultCh)
		wg.Wait()
		close(doneCh)
	}()

	go func() {
		select {
		case <-time.After(defaultTimeout):
			regionalResultCh <- RegionalResult{
				Err: bizerr.New(bizerr.ErrCodeTimeout, "multi call regional timeout, regions: %v", regions),
			}
		case <-doneCh:
			// Completed within timeout, no action needed
		}
	}()

	return regionalResultCh
}