package runner

import (
	"context"
	"errors"
	"time"

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

const (
	taskExecutionLeaseTTL      = 5 * time.Minute
	taskExecutionRenewInterval = 30 * time.Second
	leaseRenewalRetryInterval  = 2 * time.Second
)

func (s *Service) renewTaskExecutionLease(ctx context.Context, cancel context.CancelFunc, taskID string, executionToken string) {
	timer := time.NewTimer(taskExecutionRenewInterval)
	defer timer.Stop()
	leaseDeadline := time.Now().Add(taskExecutionLeaseTTL)
	for {
		select {
		case <-ctx.Done():
			return
		case <-timer.C:
			if err := s.coordinationStore.RenewTaskExecutionLease(ctx, taskID, executionToken, taskExecutionLeaseTTL); err != nil {
				if errors.Is(err, store.ErrTaskExecutionFinished) {
					return
				}
				if taskExecutionLeaseRenewalIsFatal(err, time.Now(), leaseDeadline) {
					s.logger.Warn("task execution lease lost", "taskID", taskID, "error_category", "task_execution_lease_lost", "error", err)
					cancel()
					return
				}
				s.logger.Warn("task execution lease renewal failed; retrying before lease expiry", "taskID", taskID, "error_category", "task_execution_lease_renewal_transient", "error", err)
				timer.Reset(leaseRenewalRetryInterval)
				continue
			}
			leaseDeadline = time.Now().Add(taskExecutionLeaseTTL)
			timer.Reset(taskExecutionRenewInterval)
		}
	}
}

func taskExecutionLeaseRenewalIsFatal(err error, now time.Time, leaseDeadline time.Time) bool {
	return errors.Is(err, store.ErrTaskExecutionLeaseLost) || !now.Before(leaseDeadline)
}