package runner

import (
	"context"
	"sync/atomic"
)

type submissionCancellationGate struct {
	submitted atomic.Bool
}

type submissionCancellationGateKey struct{}

func NewRequestExecutionContext(requestContext context.Context, serverContext context.Context) (context.Context, context.CancelFunc) {
	base, cancel := context.WithCancel(context.WithoutCancel(requestContext))
	gate := &submissionCancellationGate{}
	ctx := context.WithValue(base, submissionCancellationGateKey{}, gate)
	go func() {
		select {
		case <-serverContext.Done():
			cancel()
		case <-requestContext.Done():
			if !gate.submitted.Load() {
				cancel()
				return
			}
			select {
			case <-serverContext.Done():
				cancel()
			case <-base.Done():
			}
		case <-base.Done():
		}
	}()
	return ctx, cancel
}

func markUpstreamSubmissionStarted(ctx context.Context) {
	if gate, ok := ctx.Value(submissionCancellationGateKey{}).(*submissionCancellationGate); ok {
		gate.submitted.Store(true)
	}
}

func requestCancelledBeforeSubmission(ctx context.Context) bool {
	gate, ok := ctx.Value(submissionCancellationGateKey{}).(*submissionCancellationGate)
	return ok && ctx.Err() != nil && !gate.submitted.Load()
}