/*
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"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"maps"
	"sort"
	"strings"

	"github.com/LocoStack/loco-operator/api/v1alpha1"
	"github.com/LocoStack/loco-operator/pkg/templates/litellm"
	"go.yaml.in/yaml/v2"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

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

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

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

	var modelEndpoint string
	var modelName string
	toolEndpoints := make(map[string]string)
	toolNames := make([]string, 0)
	kbEndpoints := make(map[string]string)
	deps := make([]client.Object, 0)
	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)
		}
		switch dep.Kind {
		case "ExternalModel":
			fallthrough
		case "ManagedModel":
			var category string
			if extModel, ok := obj.(*v1alpha1.ExternalModel); ok {
				category = extModel.Spec.Category
				modelName = extModel.Spec.ModelName
			} else if managedModel, ok := obj.(*v1alpha1.ManagedModel); ok {
				category = managedModel.Spec.Category
				modelName = managedModel.Spec.ModelName
			} else {
				return nil, fmt.Errorf("Unsupported dependency kind for agent: %s", dep.Kind)
			}
			if category != "language" {
				return nil, fmt.Errorf("Unsupported model category for agent: %s", category)
			}
			modelEndpoint = endpoint
		case "ExternalTool":
			fallthrough
		case "ManagedTool":
			toolEndpoints[dep.Name] = endpoint
			if extTool, ok := obj.(*v1alpha1.ExternalTool); ok {
				toolNames = append(toolNames, extTool.Spec.ToolName)
			} else if managedTool, ok := obj.(*v1alpha1.ManagedTool); ok {
				toolNames = append(toolNames, managedTool.Spec.ToolName)
			}
		case "ExternalKnowledgeBase":
			fallthrough
		case "ManagedKnowledgeBase":
			kbEndpoints[dep.Name] = endpoint
		default:
			return nil, fmt.Errorf("Unsupported dependency kind for agent: %s", dep.Kind)
		}
		deps = append(deps, obj)
	}

	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)
	}

	agentConfigHash, err := r.reconcileAgentConfig(ctx, r.component.Spec.Variables["systemPrompt"], modelEndpoint, masterKey, modelName, kbEndpoints)
	if err != nil {
		return nil, fmt.Errorf("Failed to reconcile config for %s: %w", r.component.Name, err)
	}

	mcpConfigHash, err := r.reconcileMcpConfig(ctx, gatewayComp.Status.Endpoint, masterKey, toolNames, toolEndpoints)
	if err != nil {
		return nil, fmt.Errorf("Failed to reconcile MCP config for %s: %w", r.component.Name, err)
	}

	h := sha256.Sum256([]byte(agentConfigHash + mcpConfigHash))
	configHash := hex.EncodeToString(h[:])
	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
}

type modelConfig struct {
	Endpoint string `yaml:"endpoint"`
	APIKey   string `yaml:"apiKey,omitempty"`
	Name     string `yaml:"name,omitempty"`
}

type endpoint struct {
	Name     string `json:"name"`
	Endpoint string `json:"endpoint"`
}

type agentConfig struct {
	SystemPrompt   string      `yaml:"systemPrompt"`
	Model          modelConfig `yaml:"model"`
	KnowledgeBases []endpoint  `yaml:"knowledgeBases,omitempty"`
	MCPConfigPath  string      `yaml:"mcpConfigPath,omitempty"`
}

func (r *LocoAgentsReconciler) reconcileAgentConfig(ctx context.Context, systemPrompt string, modelEndpoint, modelAPIKey, modelName string, knowledgeBases map[string]string) (string, error) {
	config := agentConfig{
		SystemPrompt: systemPrompt,
		Model: modelConfig{
			Endpoint: modelEndpoint,
			APIKey:   modelAPIKey,
			Name:     modelName,
		},
		MCPConfigPath: "/etc/mcp/mcp.json",
	}
	kbNames := make([]string, 0, len(knowledgeBases))
	for kbName := range knowledgeBases {
		kbNames = append(kbNames, kbName)
	}
	sort.Strings(kbNames)
	for _, kbName := range kbNames {
		config.KnowledgeBases = append(config.KnowledgeBases, endpoint{
			Name:     kbName,
			Endpoint: knowledgeBases[kbName],
		})
	}
	yamlBytes, err := yaml.Marshal(config)
	h := sha256.Sum256(yamlBytes)
	configHash := hex.EncodeToString(h[:])
	if err != nil {
		return "", fmt.Errorf("marshaling agent config: %w", err)
	}
	secret := &corev1.Secret{ObjectMeta: metav1.ObjectMeta{Name: r.ResourceName("config"), Namespace: r.component.GetNamespace()}}
	_, err = controllerutil.CreateOrUpdate(ctx, r.client, secret, func() error {
		if secret.Labels == nil {
			secret.Labels = map[string]string{}
		}
		maps.Copy(secret.Labels, r.ResourceLabels())
		if secret.Annotations == nil {
			secret.Annotations = make(map[string]string)
		}
		secret.Annotations["hash"] = configHash
		secret.Data = map[string][]byte{
			"agent.yaml": yamlBytes,
		}
		return controllerutil.SetControllerReference(r.component, secret, r.scheme)
	})
	return configHash, err
}

type mcpConfig struct {
	MCPServers map[string]mcpServer `json:"mcpServers"`
}

type mcpServer struct {
	URL     string            `json:"url"`
	Headers map[string]string `json:"headers,omitempty"`
}

// reconcileMcpConfig writes the agent's mcp.json as a Secret (not a ConfigMap): when a
// gateway is present, the generated config embeds the resolved litellm master key value
// in the auth header, so the credential must not land in a ConfigMap.
func (r *LocoAgentsReconciler) reconcileMcpConfig(ctx context.Context, gatewayEndpoint, masterKey string, toolNames []string, toolEndpoints map[string]string) (string, error) {
	cfg := mcpConfig{MCPServers: make(map[string]mcpServer)}
	if gatewayEndpoint != "" && len(toolNames) > 0 {
		cfg.MCPServers["litellm"] = mcpServer{
			URL: fmt.Sprintf("%s/%s/mcp", gatewayEndpoint, strings.Join(toolNames, ",")),
			Headers: map[string]string{
				litellm.MCPAuthHeaderName: "Bearer " + masterKey,
			},
		}
	} else {
		for toolName, toolEndpoint := range toolEndpoints {
			cfg.MCPServers[toolName] = mcpServer{URL: toolEndpoint}
		}
	}
	jsonBytes, err := json.Marshal(cfg)
	if err != nil {
		return "", fmt.Errorf("marshaling agent mcp config: %w", err)
	}
	h := sha256.Sum256(jsonBytes)
	configHash := hex.EncodeToString(h[:])
	secret := &corev1.Secret{ObjectMeta: metav1.ObjectMeta{Name: r.ResourceName("mcp-config"), Namespace: r.component.GetNamespace()}}
	_, err = controllerutil.CreateOrUpdate(ctx, r.client, secret, func() error {
		if secret.Labels == nil {
			secret.Labels = map[string]string{}
		}
		maps.Copy(secret.Labels, r.ResourceLabels())
		if secret.Annotations == nil {
			secret.Annotations = make(map[string]string)
		}
		secret.Annotations["hash"] = configHash
		secret.Data = map[string][]byte{
			"mcp.json": jsonBytes,
		}
		return controllerutil.SetControllerReference(r.component, secret, r.scheme)
	})
	return configHash, err
}

func (r *LocoAgentsReconciler) gatewayMasterKey(ctx context.Context, gatewayComp *v1alpha1.Component) (string, error) {
	if gatewayComp == nil {
		return "", nil
	}
	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
}