package scheduler

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/constant"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/repo"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/logid"

	"go.uber.org/zap"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8slabels "k8s.io/apimachinery/pkg/labels"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

const MaxConcurrent = 10

type Reconciler struct {
	manager          manager.Manager
	auditLogTaskRepo repo.AuditLogTaskRepo
}

func NewWorkerReconciler(
	manager manager.Manager,
	auditLogTaskRepo repo.AuditLogTaskRepo,
) *Reconciler {
	return &Reconciler{
		manager:          manager,
		auditLogTaskRepo: auditLogTaskRepo,
	}
}

func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	ctx = logid.WithLogid(ctx, logid.GenLogid())
	logger := log.WithContextL(ctx).With(zap.String("worker", req.Name))

	logger.Info("start reconcile audit log worker")
	defer func() {
		logger.Info("finish reconcile audit log worker")
	}()

	worker := &batchv1.Job{}
	if err := r.manager.GetClient().Get(ctx, req.NamespacedName, worker); err != nil {
		if apierrors.IsNotFound(err) {
			logger.Info("audit log worker is not found")
			return ctrl.Result{}, nil
		}
		logger.Error("get audit worker from cache failed", zap.Error(err))
		return ctrl.Result{}, errors.New("get audit worker from cache failed")
	}

	// update any runnning task to failed when the worker is failed
	for _, cond := range worker.Status.Conditions {
		if cond.Type == batchv1.JobFailed && cond.Status == corev1.ConditionTrue {
			logger.Info("audit log worker failed, try to failed the tasks")
			// fetch task ids from worker
			taskIds := make([]string, 0)
			for _, c := range worker.Spec.Template.Spec.Containers {
				if c.Name == constant.AuditLogWorkerContainerName {
					for _, arg := range c.Args {
						if strings.HasPrefix(arg, "--task-ids=") {
							taskIds = append(taskIds, strings.TrimPrefix(arg, "--task-ids="))
						}
					}
				}
			}
			rowAffected, err := r.auditLogTaskRepo.BatchFailed(ctx, taskIds, "unknown")
			if err != nil {
				logger.Error("failed to execute audit log task batch failed", zap.Error(err))
				return ctrl.Result{}, err
			}
			logger.Info("succeed to execute audit log task batch failed", zap.Int64("row_affected", rowAffected))
		} else if cond.Type == batchv1.JobComplete && cond.Status == corev1.ConditionTrue {
			// delete the worker with background propagation policy
			return ctrl.Result{}, r.manager.GetClient().Delete(ctx, worker, client.PropagationPolicy(metav1.DeletePropagationBackground))
		}
	}
	return ctrl.Result{}, nil
}

func (r *Reconciler) SetupWithManager(mgr ctrl.Manager) error {
	c, err := controller.New("AuditLogWorker", mgr, controller.Options{
		MaxConcurrentReconciles: MaxConcurrent,
		Reconciler:              r,
	})
	if err != nil {
		return fmt.Errorf("can't create controller: %w", err)
	}

	pre, err := jobLabelSelectorPredicate(metav1.LabelSelector{
		MatchLabels: map[string]string{
			types.LabelJobType: constant.AuditLogWorkerLabelValue,
		},
	})

	if err != nil {
		return fmt.Errorf("invalid label selector: %w", err)
	}

	err = c.Watch(source.Kind(mgr.GetCache(), &batchv1.Job{}, &handler.TypedEnqueueRequestForObject[*batchv1.Job]{}, pre))
	if err != nil {
		return fmt.Errorf("can't watch audit log worker")
	}

	log.Info(context.Background(), "AuditLogWorker reconciler setup with manager successfully")
	return nil
}

// jobLabelSelectorPredicate creates a TypedPredicate for *batchv1.Job using a LabelSelector.
func jobLabelSelectorPredicate(selector metav1.LabelSelector) (predicate.TypedPredicate[*batchv1.Job], error) {
	labelSelector, err := metav1.LabelSelectorAsSelector(&selector)
	if err != nil {
		return nil, err
	}
	return predicate.NewTypedPredicateFuncs(func(job *batchv1.Job) bool {
		return labelSelector.Matches(k8slabels.Set(job.GetLabels()))
	}), nil
}