package testsupport

import (
	"context"
	"git.wamblee.org/converge/pkg/support/pprof"
	"github.com/stretchr/testify/suite"
	"log"
	"net/http"
	"os"
	_ "runtime/pprof"
	"sync"
	"time"
)

type TestFunction func() any

func RunAndWait(suite *suite.Suite, functions ...TestFunction) []any {
	wg := sync.WaitGroup{}
	wg.Add(len(functions))
	res := make([]any, len(functions))
	for i, function := range functions {
		go func() {
			defer func() {
				wg.Done()
			}()
			res[i] = function()
		}()
	}
	wg.Wait()
	return res
}

func StartPprof(port string) *http.Server {
	if os.Getenv("PPROF") == "" {
		return nil
	}
	if port == "" {
		port = ":9000"
	}
	mux := http.NewServeMux()
	pprof.RegisterPprof(mux, "/debug/pprof")
	srv := http.Server{
		Addr:    port,
		Handler: mux,
	}
	go func() {
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatalf("Could not start pprof listener: %v", err)
		}
		log.Println("Test pprof server started: " + port)
	}()
	return &srv
}

func StopPprof(ctx context.Context, server *http.Server) {
	if os.Getenv("PPROF") == "" {
		return
	}
	err := server.Shutdown(ctx)
	if err != nil {
		log.Println("Error shutting down test pprof server")
		return
	}
	log.Println("Test pprof server stopped")
}

func CreateTestContext(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
	ctx, cancelFunc := context.WithCancel(ctx)
	ctx, timeoutCancelFunc := context.WithTimeout(ctx, timeout)
	compositeCancelFunc := func() {
		timeoutCancelFunc()
		cancelFunc()
	}
	return ctx, compositeCancelFunc
}