package runner

import (
	"context"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload"
)

type workerLoadLeaseContextKey struct{}

func workerLoadLease(ctx context.Context) *workerload.Lease {
	lease, _ := ctx.Value(workerLoadLeaseContextKey{}).(*workerload.Lease)
	return lease
}

func enterWorkerWaiting(ctx context.Context) error {
	if lease := workerLoadLease(ctx); lease != nil {
		return lease.EnterWaiting()
	}
	return nil
}

func enterWorkerFinalizing(ctx context.Context) error {
	if lease := workerLoadLease(ctx); lease != nil {
		return lease.EnterFinalizing(ctx)
	}
	return nil
}