loco-operator/internal/controller/document_controller.go

/*
Copyright 2026 LocoStack.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package controller

import (
	"context"
	"fmt"
	"strings"

	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	apimeta "k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	logf "sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/LocoStack/loco-operator/api/v1alpha1"
	"github.com/LocoStack/loco-operator/internal/reconciler"
	"github.com/LocoStack/loco-operator/pkg/templates"
)

const (
	finalizerDocumentCleanup = "locostack.com/document-cleanup"
)

// DocumentReconciler reconciles a Document object
type DocumentReconciler struct {
	client.Client
	Scheme *runtime.Scheme
}

// +kubebuilder:rbac:groups=locostack.com,resources=documents,verbs=get;list;watch;create;update;patch;delete
// +kubebuilder:rbac:groups=locostack.com,resources=documents/status,verbs=get;update;patch
// +kubebuilder:rbac:groups=locostack.com,resources=documents/finalizers,verbs=update
// +kubebuilder:rbac:groups=locostack.com,resources=knowledgebases,verbs=get;list;watch
// +kubebuilder:rbac:groups=locostack.com,resources=components,verbs=get;list;watch
// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;update;patch;delete

func (r *DocumentReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := logf.FromContext(ctx)

	doc := &v1alpha1.Document{}
	if err := r.Get(ctx, req.NamespacedName, doc); err != nil {
		if apierrors.IsNotFound(err) {
			return ctrl.Result{}, nil
		}
		log.Error(err, "Failed to get Document", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
		return ctrl.Result{}, err
	}

	patch := client.MergeFrom(doc.DeepCopy())

	ingestionJobTemplate, egestionJobTemplate, stack, endpoint, err := r.reconcileKB(ctx, doc)
	if err != nil {
		return ctrl.Result{}, err
	}

	if !controllerutil.ContainsFinalizer(doc, finalizerDocumentCleanup) {
		patch := client.MergeFrom(doc.DeepCopy())
		controllerutil.AddFinalizer(doc, finalizerDocumentCleanup)
		return ctrl.Result{}, r.Patch(ctx, doc, patch)
	}

	if !doc.DeletionTimestamp.IsZero() {
		egested, egestionFailed, err := r.reconcileEgestion(ctx, doc, egestionJobTemplate, stack, endpoint)
		if err != nil {
			return ctrl.Result{}, err
		}
		if !egested {
			return ctrl.Result{}, nil
		}
		if egestionFailed {
			// Handle egestion failure if necessary
			return ctrl.Result{}, fmt.Errorf("Egestion failed for document %s/%s", req.NamespacedName.Namespace, req.NamespacedName.Name)
		}
		controllerutil.RemoveFinalizer(doc, finalizerDocumentCleanup)
		return ctrl.Result{}, r.Patch(ctx, doc, patch)
	}

	ingested, ingestionFailed, err := r.reconcileIngestion(ctx, doc, ingestionJobTemplate, stack, endpoint)
	if err != nil {
		return ctrl.Result{}, err
	}

	if ingested {
		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
			Type:               "Ready",
			Status:             metav1.ConditionTrue,
			Reason:             "Ingested",
			Message:            "Document has been ingested into the KnowledgeBase",
			ObservedGeneration: doc.Generation,
		})
	} else {
		if ingestionFailed {
			apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
				Type:               "Ready",
				Status:             metav1.ConditionFalse,
				Reason:             "Failed",
				Message:            "Document failed to be ingested into the KnowledgeBase",
				ObservedGeneration: doc.Generation,
			})
		} else {
			apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
				Type:               "Ready",
				Status:             metav1.ConditionFalse,
				Reason:             "Ingesting",
				Message:            "Document is being ingested into the KnowledgeBase",
				ObservedGeneration: doc.Generation,
			})
		}

	}

	doc.Status.ObservedGeneration = doc.Generation
	if err := r.Status().Patch(ctx, doc, patch); err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	log.Info("Reconciled Document", "namespace", req.NamespacedName.Namespace, "name", req.NamespacedName.Name)
	return ctrl.Result{}, nil
}

func (r *DocumentReconciler) reconcileKB(ctx context.Context, doc *v1alpha1.Document) (*v1alpha1.Template, *v1alpha1.Template, *v1alpha1.Stack, string, error) {
	log := logf.FromContext(ctx)
	kbRef := doc.Spec.KnowledgeBaseRef
	if kbRef == nil {
		// TODO
		return nil, nil, nil, "", fmt.Errorf("KnowledgeBaseRef is required in Document spec")
	}
	var stackRef *corev1.LocalObjectReference
	var ingestionJobTemplate *v1alpha1.Template
	var egestionJobTemplate *v1alpha1.Template
	var endpoint string
	if kbRef.Kind == "ExternalKnowledgeBase" {
		// TODO
	} else if kbRef.Kind == "ManagedKnowledgeBase" {
		kb := &v1alpha1.ManagedKnowledgeBase{}
		if err := r.Get(ctx, client.ObjectKey{Namespace: doc.Namespace, Name: doc.Spec.KnowledgeBaseRef.Name}, kb); err != nil {
			if apierrors.IsNotFound(err) {
				apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
					Type:               "KnowledgeBaseAvailable",
					Status:             metav1.ConditionFalse,
					Reason:             "KnowledgeBaseNotFound",
					Message:            "Referenced KnowledgeBase not found",
					ObservedGeneration: doc.Generation,
				})
				return nil, nil, nil, "", nil
			}
			log.Error(err, "Failed to get referenced ManagedKnowledgeBase", "namespace", doc.Namespace, "name", doc.Spec.KnowledgeBaseRef.Name)
			return nil, nil, nil, "", err
		}
		stackRef = kb.Spec.StackRef
		ingestionJobTemplate = kb.Spec.IngestionJob
		egestionJobTemplate = kb.Spec.EgestionJob
		endpoint = kb.Status.Endpoint
	}
	stack := &v1alpha1.Stack{}
	if err := r.Get(ctx, client.ObjectKey{Namespace: doc.Namespace, Name: stackRef.Name}, stack); err != nil {
		if apierrors.IsNotFound(err) {
			apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
				Type:               "KnowledgeBaseAvailable",
				Status:             metav1.ConditionFalse,
				Reason:             "StackNotFound",
				Message:            "Stack for referenced KnowledgeBase not found",
				ObservedGeneration: doc.Generation,
			})
			return nil, nil, nil, "", nil
		}
		log.Error(err, "Failed to get Stack for referenced ManagedKnowledgeBase", "namespace", doc.Namespace, "name", stackRef.Name)
		return nil, nil, nil, "", err
	}
	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
		Type:               "KnowledgeBaseAvailable",
		Status:             metav1.ConditionTrue,
		Reason:             "KnowledgeBaseAvailable",
		Message:            "Referenced KnowledgeBase is available",
		ObservedGeneration: doc.Generation,
	})
	return ingestionJobTemplate, egestionJobTemplate, stack, endpoint, nil
}

func (r *DocumentReconciler) reconcileArtifacts(ctx context.Context, stack *v1alpha1.Stack, doc *v1alpha1.Document) (*corev1.PersistentVolumeClaim, bool, bool, error) {
	artifactsReconciler := reconciler.NewArtifactReconciler(r.Client, r.Scheme, stack, "document", doc, &doc.Spec.Source)
	return artifactsReconciler.ReconcileArtifact(ctx)
}

func (r *DocumentReconciler) reconcileIngestion(ctx context.Context, doc *v1alpha1.Document, jobTmpl *v1alpha1.Template, stack *v1alpha1.Stack, kbEndpoint string) (bool, bool, error) {
	if doc.Status.Phase == "Ready" || doc.Status.Phase == "Failed" {
		return true, doc.Status.Phase == "Failed", nil
	}

	tmpl, err := templates.Manager.ResolveTemplate(jobTmpl, "")
	if err != nil {
		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
			Type:               "IngestionJobResolved",
			Status:             metav1.ConditionFalse,
			Reason:             "TemplateResolutionFailed",
			Message:            "Failed to resolve ingestion job template: " + err.Error(),
			ObservedGeneration: doc.Generation,
		})
		return false, true, err
	}
	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
		Type:               "IngestionJobResolved",
		Status:             metav1.ConditionTrue,
		Reason:             "TemplateResolved",
		Message:            "Ingestion job template resolved successfully",
		ObservedGeneration: doc.Generation,
	})

	sourcePvc, sourceReady, sourceFailed, err := r.reconcileArtifacts(ctx, stack, doc)
	if err != nil {
		return false, true, err
	}
	if sourceFailed {
		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
			Type:               "SourceReady",
			Status:             metav1.ConditionFalse,
			Reason:             "SourceFailed",
			Message:            "Document source is not available",
			ObservedGeneration: doc.Generation,
		})
		return false, true, nil
	} else if !sourceReady {
		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
			Type:               "SourceReady",
			Status:             metav1.ConditionFalse,
			Reason:             "SourceNotReady",
			Message:            "Document source is not ready",
			ObservedGeneration: doc.Generation,
		})
		return false, false, nil
	}
	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
		Type:               "SourceReady",
		Status:             metav1.ConditionTrue,
		Reason:             "SourceReady",
		Message:            "Document source is ready",
		ObservedGeneration: doc.Generation,
	})

	doc.Status.Phase = "Ingesting"
	jobReconciler := reconciler.NewJobReconciler(r.Client, r.Scheme, stack, "DocumentIngestion", doc.Name, doc)
	variables := map[string]string{
		"spec.runtime.pvc":                   sourcePvc.Name,
		"spec.docId":                         doc.Name,
		"spec.chunkingStrategy.strategy":     doc.Spec.ChunkingStrategy.Strategy,
		"spec.chunkingStrategy.chunkSize":    fmt.Sprint(doc.Spec.ChunkingStrategy.ChunkSize),
		"spec.chunkingStrategy.chunkOverlap": fmt.Sprint(doc.Spec.ChunkingStrategy.ChunkOverlap),
		"spec.chunkingStrategy.separators":   strings.Join(doc.Spec.ChunkingStrategy.Separators, ","),
		"spec.kbEndpoint":                    kbEndpoint,
	}
	if doc.Spec.Source.PVC != nil {
		variables["spec.filePath"] = doc.Spec.Source.PVC.FilePath
	}
	done, failed, err := jobReconciler.ReconcileJob(ctx, &tmpl, variables)
	if done {
		doc.Status.Phase = "Ready"
	} else if failed {
		doc.Status.Phase = "Failed"
	}
	return done, failed, err
}

func (r *DocumentReconciler) reconcileEgestion(ctx context.Context, doc *v1alpha1.Document, jobTmpl *v1alpha1.Template, stack *v1alpha1.Stack, kbEndpoint string) (bool, bool, error) {
	if doc.Status.Phase != "Ready" {
		// If the document is not in the "Ready" phase, we do not need egestion.
		return true, false, nil
	}

	doc.Status.Phase = "Egesting"
	jobReconciler := reconciler.NewJobReconciler(r.Client, r.Scheme, stack, "DocumentEgestion", doc.Name, doc)
	tmpl, err := templates.Manager.ResolveTemplate(jobTmpl, "")
	if err != nil {
		apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
			Type:               "EgestionJobResolved",
			Status:             metav1.ConditionFalse,
			Reason:             "TemplateResolutionFailed",
			Message:            "Failed to resolve egestion job template: " + err.Error(),
			ObservedGeneration: doc.Generation,
		})
		return false, true, err
	}
	apimeta.SetStatusCondition(&doc.Status.Conditions, metav1.Condition{
		Type:               "EgestionJobResolved",
		Status:             metav1.ConditionTrue,
		Reason:             "TemplateResolved",
		Message:            "Egestion job template resolved successfully",
		ObservedGeneration: doc.Generation,
	})
	variables := map[string]string{
		"spec.docId":      doc.Name,
		"spec.kbEndpoint": kbEndpoint,
	}
	return jobReconciler.ReconcileJob(ctx, &tmpl, variables)
}

// SetupWithManager sets up the controller with the Manager.
func (r *DocumentReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&v1alpha1.Document{}).
		Owns(&batchv1.Job{}).
		Owns(&corev1.PersistentVolumeClaim{}).
		Named("document").
		Complete(r)
}