package activator

import (
	"context"
	"sync"
	"time"
)

type cleanupBudgetKey struct{}
type cleanupContextKey struct{}
type cleanupBudget struct {
	mu     sync.Mutex
	base   context.Context
	source context.Context
	parent *cleanupBudget
	limit  time.Duration
	ctx    context.Context
	cancel context.CancelFunc
}

func (b *cleanupBudget) start() context.Context {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.ctx == nil {
		base := b.base
		if b.parent != nil && b.parent.source.Err() != nil {
			base = b.parent.start()
		}
		b.ctx, b.cancel = context.WithTimeout(base, b.limit)
	}
	return b.ctx
}

// WithCleanupBudget shares the cancellation grace period across a request,
// its nested failure paths, and the owner's final join. Normal work is untimed.
func WithCleanupBudget(ctx context.Context, limit time.Duration) (context.Context, context.CancelFunc) {
	parent, _ := ctx.Value(cleanupBudgetKey{}).(*cleanupBudget)
	b := &cleanupBudget{base: context.WithoutCancel(ctx), source: ctx, parent: parent, limit: limit}
	stop := context.AfterFunc(ctx, func() { b.start() })
	return context.WithValue(ctx, cleanupBudgetKey{}, b), func() { stop(); b.start(); b.cancel() }
}

// CleanupContext keeps an existing cleanup deadline. Cancellation starts the
// owner's single grace period; it never grants each nested call a fresh one.
func CleanupContext(ctx context.Context, limit time.Duration) (context.Context, context.CancelFunc) {
	if ctx.Value(cleanupContextKey{}) != nil {
		return context.WithTimeout(ctx, limit)
	}
	base := context.WithoutCancel(ctx)
	var stop func() bool
	if b, ok := ctx.Value(cleanupBudgetKey{}).(*cleanupBudget); ok {
		if ctx.Err() != nil {
			base = b.start()
		}
		cleanup, cancel := context.WithTimeout(base, limit)
		stop = context.AfterFunc(ctx, func() { budget := b.start(); context.AfterFunc(budget, cancel) })
		return context.WithValue(cleanup, cleanupContextKey{}, true), func() { stop(); cancel() }
	}
	// Callers without an owner budget still cannot extend an explicit deadline.
	if deadline, ok := ctx.Deadline(); ok {
		bounded, cancel := context.WithDeadline(base, deadline)
		cleanup, done := context.WithTimeout(bounded, limit)
		return context.WithValue(cleanup, cleanupContextKey{}, true), func() { done(); cancel() }
	}
	cleanup, cancel := context.WithTimeout(base, limit)
	return context.WithValue(cleanup, cleanupContextKey{}, true), cancel
}