package jobs

import (
	"context"
	"fmt"
	"time"

	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/klog/v2"
)

// WaitForStopped waits for a Job to stop running (= completed or failed)
func WaitForStopped(ctx context.Context, kubeClient kubernetes.Interface, jobName string, jobNamespace string, timeout time.Duration) error {
	klog.Infof("Waiting for job %s to complete or fail (=not running anymore) (timeout: %v)", jobName, timeout)
	isStopped := func(job batchv1.Job) (bool, error) {
		return IsStopped(job), nil
	}
	return waitWithConditionFunc(ctx, kubeClient, jobName, jobNamespace, timeout, isStopped)
}

// WaitForCompletion waits for a Job to complete (= succeed)
func WaitForCompletion(ctx context.Context, kubeClient kubernetes.Interface, jobName string, jobNamespace string, timeout time.Duration) error {
	klog.Infof("Waiting for job %s to complete (=succeed) (timeout: %v)", jobName, timeout)
	isComplete := func(job batchv1.Job) (bool, error) {
		if IsFailed(job) {
			// return error to signal that we don't need to wait any longer
			return false, fmt.Errorf("job %s failed", jobName)
		}
		return IsComplete(job), nil
	}
	return waitWithConditionFunc(ctx, kubeClient, jobName, jobNamespace, timeout, isComplete)
}

// waitWithConditionFunc waits for a Job to fulfill given conditionFunc
func waitWithConditionFunc(ctx context.Context, kubeClient kubernetes.Interface, jobName string, jobNamespace string, timeout time.Duration, conditionFunc func(job batchv1.Job) (bool, error)) error {
	return wait.PollUntilContextTimeout(ctx, 5*time.Second, timeout, true, func(ctx context.Context) (bool, error) {
		job, err := kubeClient.BatchV1().Jobs(jobNamespace).Get(ctx, jobName, v1.GetOptions{})
		if err != nil {
			// Ignore errors (including NotFound) to avoid returning early. The job might be
			// temporarily missing during deletion/recreation cycles.
			klog.Warningf("Failed to get job %s: %v", jobName, err)
			return false, nil
		}

		// Check if job condition is fulfilled
		if isFulfilled, err := conditionFunc(*job); err != nil {
			klog.Infof("Job %s condition can not be fulfilled anymore", jobName)
			return false, err
		} else if isFulfilled {
			klog.Infof("Job %s condition fulfilled", jobName)
			return true, nil
		}

		return false, nil
	})
}

// DeleteAndWait deletes a job and waits until it disappears from the API.
// Uses Background propagation to automatically delete child pods.
func DeleteAndWait(ctx context.Context, kubeClient kubernetes.Interface, jobName string, jobNamespace string) error {
	klog.V(4).Infof("deleting oldJob %s", jobName)
	oldJob, err := kubeClient.BatchV1().Jobs(jobNamespace).Get(ctx, jobName, v1.GetOptions{})
	if err != nil {
		if apierrors.IsNotFound(err) {
			klog.V(4).Infof("job %s didn't exist", jobName)
			return nil
		}
		return err
	}
	oldJobUID := oldJob.GetUID()

	// Delete job with Background propagation to automatically clean up pods
	propagationPolicy := v1.DeletePropagationBackground
	err = kubeClient.BatchV1().Jobs(jobNamespace).Delete(ctx, jobName, v1.DeleteOptions{
		PropagationPolicy: &propagationPolicy,
	})
	if err != nil {
		if apierrors.IsNotFound(err) {
			return nil
		}
		return err
	}

	// wait until the oldJob is deleted
	klog.V(4).Infof("waiting for job %s to be deleted", jobName)
	return wait.PollUntilContextTimeout(ctx, 5*time.Second, 1*time.Minute, true, func(ctx context.Context) (bool, error) {
		newJob, err := kubeClient.BatchV1().Jobs(jobNamespace).Get(ctx, jobName, v1.GetOptions{})
		if apierrors.IsNotFound(err) {
			return true, nil
		}
		if err != nil {
			klog.Warningf("error checking job %s deletion status: %v", jobName, err)
			return false, nil // retry on transient errors
		}
		// job might be recreated already, check UID
		return newJob.GetUID() != oldJobUID, nil
	})
}

func IsStopped(job batchv1.Job) bool {
	// Check for explicit conditions
	if IsComplete(job) || IsFailed(job) {
		return true
	}

	// Check for FailureTarget condition (Kubernetes 1.31+)
	// FailureTarget means job is targeting failure but pods may still be terminating
	if IsConditionTrue(job.Status.Conditions, batchv1.JobConditionType("FailureTarget")) {
		klog.V(2).Infof("Job %s considered stopped: FailureTarget condition is True", job.Name)
		return true
	}

	// Also check if job has no active pods and has exceeded backoff limit
	// This handles cases where pods are stuck (e.g., Terminating on dead node) and
	// Kubernetes hasn't set the Failed condition yet
	if job.Status.Active == 0 && job.Status.Failed > 0 {
		backoffLimit := int32(6) // Kubernetes default
		if job.Spec.BackoffLimit != nil {
			backoffLimit = *job.Spec.BackoffLimit
		}
		if job.Status.Failed > backoffLimit {
			klog.V(2).Infof("Job %s considered stopped: Active=0, Failed=%d > BackoffLimit=%d (no explicit condition set)",
				job.Name, job.Status.Failed, backoffLimit)
			return true
		}
	}

	return false
}

func IsComplete(job batchv1.Job) bool {
	return IsConditionTrue(job.Status.Conditions, batchv1.JobComplete)
}

func IsFailed(job batchv1.Job) bool {
	return IsConditionTrue(job.Status.Conditions, batchv1.JobFailed)
}

func IsConditionTrue(conditions []batchv1.JobCondition, conditionType batchv1.JobConditionType) bool {
	return IsConditionPresentAndEqual(conditions, conditionType, corev1.ConditionTrue)
}

func IsConditionPresentAndEqual(conditions []batchv1.JobCondition, conditionType batchv1.JobConditionType, status corev1.ConditionStatus) bool {
	for _, condition := range conditions {
		if condition.Type == conditionType {
			return condition.Status == status
		}
	}
	return false
}
