/*
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 reconciler

import (
	"context"
	"fmt"
	"maps"
	"sort"

	"github.com/LocoStack/loco-operator/api/v1alpha1"
	"github.com/LocoStack/loco-operator/pkg/templates/litellm"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type LocoContextReconciler struct {
	*DefaultComponentReconciler
	client    client.Client
	scheme    *runtime.Scheme
	stack     *v1alpha1.Stack
	component *v1alpha1.Component
}

func NewLocoContextReconciler(client client.Client, scheme *runtime.Scheme, stack *v1alpha1.Stack, component *v1alpha1.Component) *LocoContextReconciler {
	return &LocoContextReconciler{
		DefaultComponentReconciler: NewDefaultComponentReconciler(client, scheme, stack, component),
		client:                     client,
		scheme:                     scheme,
		stack:                      stack,
		component:                  component,
	}
}

func (r *LocoContextReconciler) ReconcileComponent(ctx context.Context, tmpl *v1alpha1.Template, variables map[string]string) ([]client.Object, error) {
	endpointResolver := NewEndpointResolver(r.client, r.scheme, r.stack.Namespace)

	deps := make([]client.Object, 0)
	var embeddingModelEndpoint string
	var embeddingModelName string
	var rerankerModelName string
	var rerankerModelEndpoint string
	for _, dep := range r.component.Spec.Dependencies {
		obj, endpoint, err := endpointResolver.ResolveEndpoint(ctx, r.stack, dep)
		if err != nil {
			return nil, fmt.Errorf("Failed to resolve endpoint for %s %s: %w", dep.Kind, dep.Name, err)
		}
		var category string
		switch dep.Kind {
		case "ExternalModel":
			category = obj.(*v1alpha1.ExternalModel).Spec.Category
			switch category {
			case "embedding":
				embeddingModelEndpoint = endpoint
				embeddingModelName = obj.(*v1alpha1.ExternalModel).Spec.ModelName
			case "reranker":
				rerankerModelEndpoint = endpoint
				rerankerModelName = obj.(*v1alpha1.ExternalModel).Spec.ModelName
			}
		case "ManagedModel":
			category = obj.(*v1alpha1.ManagedModel).Spec.Category
			switch category {
			case "embedding":
				embeddingModelEndpoint = endpoint
				embeddingModelName = obj.(*v1alpha1.ManagedModel).Spec.ModelName
			case "reranker":
				rerankerModelEndpoint = endpoint
				rerankerModelName = obj.(*v1alpha1.ManagedModel).Spec.ModelName
			}
		}

		deps = append(deps, obj)
	}

	var vectorStoreEndpoint string
	var graphStoreEndpoint string
	if r.stack != nil {
		vectorStore := r.stack.Status.ComponentStatus.VectorStore
		if vectorStore != "" && vectorStore != "disabled" {
			if obj, endpoint, err := r.resolveComponentEndpoint(ctx, r.stack, "VectorStore"); err != nil {
				return nil, fmt.Errorf("Failed to resolve vector store endpoint: %w", err)
			} else if obj != nil {
				deps = append(deps, obj)
				vectorStoreEndpoint = endpoint
			}
		}
		graphStore := r.stack.Status.ComponentStatus.GraphStore
		if graphStore != "" && graphStore != "disabled" {
			if obj, endpoint, err := r.resolveComponentEndpoint(ctx, r.stack, "GraphStore"); err != nil {
				return nil, fmt.Errorf("Failed to resolve graph store endpoint: %w", err)
			} else if obj != nil {
				deps = append(deps, obj)
				graphStoreEndpoint = endpoint
			}
		}
	}

	gatewayComp, err := endpointResolver.resolveGatewayEndpoint(ctx, r.stack)
	if err != nil {
		return nil, fmt.Errorf("Failed to resolve gateway for %s: %w", r.component.Name, err)
	}
	masterKey, err := r.gatewayMasterKey(ctx, gatewayComp)
	if err != nil {
		return nil, fmt.Errorf("resolving gateway master key: %w", err)
	}

	// sort deps to make sure the order is deterministic
	deps = sortObjectsByName(deps)

	config := map[string]string{
		"embedding_model_endpoint": embeddingModelEndpoint,
		"embedding_model_api_key":  masterKey,
		"embedding_model_name":     embeddingModelName,
		"reranker_model_endpoint":  rerankerModelEndpoint,
		"reranker_model_api_key":   masterKey,
		"reranker_model_name":      rerankerModelName,
		"vector_store_endpoint":    vectorStoreEndpoint,
		"graph_store_endpoint":     graphStoreEndpoint,
	}
	configHash, err := r.ReconcileConfig(ctx, "config.yaml", config, true)
	if err != nil {
		return nil, fmt.Errorf("Failed to reconcile config for %s: %w", r.component.Name, err)
	}

	if tmpl.Metadata.Annotations == nil {
		tmpl.Metadata.Annotations = make(map[string]string)
	}
	maps.Copy(tmpl.Metadata.Annotations, map[string]string{
		"locostack.com/configHash": configHash,
	})

	if _, err := r.DefaultComponentReconciler.ReconcileComponent(ctx, tmpl, variables); err != nil {
		return nil, err
	}
	return deps, nil
}

func (r *LocoContextReconciler) resolveComponentEndpoint(ctx context.Context, stack *v1alpha1.Stack, component string) (client.Object, string, error) {
	compList := &v1alpha1.ComponentList{}
	err := r.client.List(ctx, compList, client.InNamespace(stack.Namespace), client.MatchingLabels{"locostack.com/stack": stack.Name, "locostack.com/component": component})
	if err != nil {
		return nil, "", fmt.Errorf("failed to list gateway components: %v", err)
	}
	if len(compList.Items) == 0 {
		return nil, "", nil
	}
	endpointResolver := NewEndpointResolver(r.client, r.scheme, stack.Namespace)
	return endpointResolver.ResolveEndpoint(ctx, stack, corev1.ObjectReference{Kind: "Component", Name: compList.Items[0].Name})
}

func (r *LocoContextReconciler) gatewayMasterKey(ctx context.Context, gatewayComp *v1alpha1.Component) (string, error) {
	secret := &corev1.Secret{}
	secretName := litellm.MasterKeySecretName(gatewayComp.Name)
	if err := r.client.Get(ctx, client.ObjectKey{Name: secretName, Namespace: r.component.GetNamespace()}, secret); err != nil {
		return "", fmt.Errorf("Failed to get gateway master key secret %s: %w", secretName, err)
	}
	key, ok := secret.Data[litellm.LITELLM_AUTH_SECRET_KEY]
	if !ok {
		return "", fmt.Errorf("Secret %s does not contain key %q", secretName, litellm.LITELLM_AUTH_SECRET_KEY)
	}
	return string(key), nil
}

func sortObjectsByName(objs []client.Object) []client.Object {
	sorted := make([]client.Object, len(objs))
	copy(sorted, objs)
	// sort by namespace/name
	sort.Slice(sorted, func(i, j int) bool {
		return sorted[i].GetNamespace()+"/"+sorted[i].GetName() < sorted[j].GetNamespace()+"/"+sorted[j].GetName()
	})
	return sorted
}