package subprocess

import (
	"context"
	"errors"
	"fmt"
	"os/exec"
	"time"
)

const (
	gracefulTerminationWait = 2 * time.Second
	forcefulTerminationWait = 2 * time.Second
)

// ownedProcessTree owns every process started by a command invocation.
// Implementations must tolerate a leader that has already exited.
type ownedProcessTree interface {
	Start(*exec.Cmd) error
	TerminateGracefully() error
	TerminateForcefully() error
	Dispose() error
}

func waitForOwnedCommand(ctx context.Context, tree ownedProcessTree, waitCh <-chan error, captureLimits <-chan *captureLimitError) (waitErr, ctxErr error, captureLimit *captureLimitError, cleanupErr error) {
	select {
	case waitErr = <-waitCh:
		return waitErr, nil, nil, nil
	case <-ctx.Done():
		ctxErr = ctx.Err()
	case captureLimit = <-captureLimits:
	}

	cleanupErr = tree.TerminateGracefully()
	gracefulTimer := time.NewTimer(gracefulTerminationWait)
	defer gracefulTimer.Stop()

	select {
	case waitErr = <-waitCh:
		// The leader may exit before descendants finish graceful shutdown.
		cleanupErr = joinErrors(cleanupErr, tree.TerminateForcefully())
		return waitErr, ctxErr, captureLimit, cleanupErr
	case <-gracefulTimer.C:
	}

	cleanupErr = joinErrors(cleanupErr, tree.TerminateForcefully())
	forcefulTimer := time.NewTimer(forcefulTerminationWait)
	defer forcefulTimer.Stop()

	select {
	case waitErr = <-waitCh:
		return waitErr, ctxErr, captureLimit, cleanupErr
	case <-forcefulTimer.C:
		return nil, ctxErr, captureLimit, joinErrors(cleanupErr, fmt.Errorf("owned subprocess did not reap within %s after forceful termination", forcefulTerminationWait))
	}
}

func joinErrors(errs ...error) error {
	filtered := make([]error, 0, len(errs))
	for _, err := range errs {
		if err != nil {
			filtered = append(filtered, err)
		}
	}
	return errors.Join(filtered...)
}