package httpz

import (
	"context"
	"net"
	"net/http"
	"sync"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/syncz"
)

func RunServer(ctx context.Context, srv *http.Server, listener net.Listener, listenerGracePeriod, shutdownTimeout time.Duration) error {
	var wg sync.WaitGroup
	defer wg.Wait() // wait for goroutine to shut down active connections
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	wg.Add(1)
	context.AfterFunc(ctx, func() { //nolint: contextcheck
		defer wg.Done()
		time.Sleep(listenerGracePeriod)
		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
		defer shutdownCancel()
		if srv.Shutdown(shutdownCtx) != nil {
			srv.Close() //nolint: errcheck,gosec
			// unhandled error above, but we are terminating anyway
		}
	})

	err := srv.Serve(listener)
	if err == http.ErrServerClosed {
		// Clean shutdown
		return nil
	}
	// Failed to start or dirty shutdown
	return err
}

// RunServerWithUpgradeSupport runs the HTTP server, tracking lifetimes of all hijacked connections.
// http.Server#Close() and http.Server#Shutdown() do not shut down hijacked connections. This helper should be used
// to run servers that support HTTP connection upgrade.
// Note: this function does not unblock read or write calls on the underlying network connection,
// it only cancels the connection's context. Calling code is responsible for reacting to context cancellation, potentially
// by calling Close() on the hijacked network connection.
func RunServerWithUpgradeSupport(ctx context.Context, srv *http.Server, listener net.Listener, listenerGracePeriod, shutdownTimeout time.Duration) error {
	// BaseContext is called from the same goroutine in Serve() so don't need locking around this var.
	// cancels may be empty if srv.Server() returns with an error before calling BaseContext().
	// We use a slice to represent "zero or more" vs a cancel var that can be nil.
	var cancels []context.CancelFunc
	shutdownStart := make(chan time.Time, 1)

	wg := syncz.NewWaitGroup()
	defer func() {
		// At this stage we know RunServer() has returned so all non-upgraded connections had been closed.
		timeLeft := shutdownTimeout - time.Since(<-shutdownStart)
		if wg.Wait(timeLeft) {
			// No in-flight upgraded connections left.
			return
		}
		// Cancel contexts of any in-flight upgraded connections.
		for _, cancel := range cancels {
			cancel()
		}
		// Wait for them to wrap up. This should return almost immediately.
		wg.Wait(time.Hour)
	}()

	bc := srv.BaseContext
	if bc == nil {
		bc = func(_ net.Listener) context.Context {
			return context.Background()
		}
	}
	srv.BaseContext = func(l net.Listener) context.Context {
		bcCtx := bc(l)
		bcCtx, cancel := context.WithCancel(bcCtx)
		cancels = append(cancels, cancel)
		return bcCtx
	}
	h := srv.Handler
	srv.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		wg.Add()
		defer wg.Done()
		h.ServeHTTP(w, r)
	})
	srv.RegisterOnShutdown(func() {
		shutdownStart <- time.Now()
	})
	return RunServer(ctx, srv, listener, listenerGracePeriod, shutdownTimeout)
}
