// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package virtualmcp

import (
	"fmt"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"

	mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
	thvjson "github.com/stacklok/toolhive/pkg/json"
	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
	"github.com/stacklok/toolhive/test/e2e/images"
)

var _ = Describe("VirtualMCPServer Composite Parallel Workflow", Ordered, func() {
	var (
		testNamespace   = "default"
		mcpGroupName    = "test-composite-par-group"
		vmcpServerName  = "test-vmcp-composite-par"
		backend1Name    = "yardstick-par-a"
		backend2Name    = "yardstick-par-b"
		timeout         = 3 * time.Minute
		pollingInterval = 1 * time.Second
		vmcpNodePort    int32

		// Composite tool name
		compositeToolName = "parallel_echo"
	)

	BeforeAll(func() {
		By("Creating MCPGroup for composite parallel test")
		CreateMCPGroupAndWait(ctx, k8sClient, mcpGroupName, testNamespace,
			"Test MCP Group for composite parallel E2E tests", timeout, pollingInterval)

		By("Creating yardstick backend MCPServers in parallel")
		CreateMultipleMCPServersInParallel(ctx, k8sClient, []BackendConfig{
			{Name: backend1Name, Namespace: testNamespace, GroupRef: mcpGroupName, Image: images.YardstickServerImage},
			{Name: backend2Name, Namespace: testNamespace, GroupRef: mcpGroupName, Image: images.YardstickServerImage},
		}, timeout, pollingInterval)

		By("Creating VirtualMCPServer with composite parallel workflow")
		vmcpServer := &mcpv1beta1.VirtualMCPServer{
			ObjectMeta: metav1.ObjectMeta{
				Name:      vmcpServerName,
				Namespace: testNamespace,
			},
			Spec: mcpv1beta1.VirtualMCPServerSpec{
				GroupRef: &mcpv1beta1.MCPGroupRef{Name: mcpGroupName},
				Config: vmcpconfig.Config{
					Group: mcpGroupName,
					Aggregation: &vmcpconfig.AggregationConfig{
						ConflictResolution: "prefix",
					},
					// Define a composite tool that echoes to both backends in parallel
					// Steps without DependsOn can execute concurrently
					CompositeTools: []vmcpconfig.CompositeToolConfig{
						{
							Name:        compositeToolName,
							Description: "Echoes message to both backends in parallel, then combines results",
							Parameters: thvjson.NewMap(map[string]any{
								"type": "object",
								"properties": map[string]any{
									"message": map[string]any{
										"type":        "string",
										"description": "The message to echo in parallel to both backends",
									},
								},
								"required": []any{"message"},
							}),
							Timeout: vmcpconfig.Duration(60 * time.Second),
							Steps: []vmcpconfig.WorkflowStepConfig{
								{
									// Step 1: Echo to backend1 (no dependencies - runs in parallel)
									ID:   "echo_backend1",
									Type: "tool",
									Tool: fmt.Sprintf("%s.echo", backend1Name),
									Arguments: thvjson.NewMap(map[string]any{
										"input": "backend1: {{ .params.message }}",
									}),
								},
								{
									// Step 2: Echo to backend2 (no dependencies - runs in parallel with step 1)
									ID:   "echo_backend2",
									Type: "tool",
									Tool: fmt.Sprintf("%s.echo", backend2Name),
									Arguments: thvjson.NewMap(map[string]any{
										"input": "backend2: {{ .params.message }}",
									}),
								},
								{
									// Step 3: Final aggregation - depends on both parallel steps
									ID:        "combine_results",
									Type:      "tool",
									Tool:      fmt.Sprintf("%s.echo", backend1Name),
									DependsOn: []string{"echo_backend1", "echo_backend2"},
									Arguments: thvjson.NewMap(map[string]any{
										"input": "Combined: [{{ .steps.echo_backend1.result }}] + [{{ .steps.echo_backend2.result }}]",
									}),
								},
							},
						},
					},
				},
				IncomingAuth: &mcpv1beta1.IncomingAuthConfig{
					Type: "anonymous",
				},
				ServiceType: "NodePort",
			},
		}
		Expect(k8sClient.Create(ctx, vmcpServer)).To(Succeed())

		By("Waiting for VirtualMCPServer to be ready")
		WaitForVirtualMCPServerReady(ctx, k8sClient, vmcpServerName, testNamespace, timeout, pollingInterval)

		By("Getting NodePort for VirtualMCPServer")
		vmcpNodePort = GetVMCPNodePort(ctx, k8sClient, vmcpServerName, testNamespace, timeout, pollingInterval)

		By(fmt.Sprintf("VirtualMCPServer accessible at http://localhost:%d", vmcpNodePort))
	})

	AfterAll(func() {
		By("Cleaning up VirtualMCPServer")
		vmcpServer := &mcpv1beta1.VirtualMCPServer{
			ObjectMeta: metav1.ObjectMeta{
				Name:      vmcpServerName,
				Namespace: testNamespace,
			},
		}
		_ = k8sClient.Delete(ctx, vmcpServer)

		By("Cleaning up backend MCPServers")
		for _, backendName := range []string{backend1Name, backend2Name} {
			backend := &mcpv1beta1.MCPServer{
				ObjectMeta: metav1.ObjectMeta{
					Name:      backendName,
					Namespace: testNamespace,
				},
			}
			_ = k8sClient.Delete(ctx, backend)
		}

		By("Cleaning up MCPGroup")
		mcpGroup := &mcpv1beta1.MCPGroup{
			ObjectMeta: metav1.ObjectMeta{
				Name:      mcpGroupName,
				Namespace: testNamespace,
			},
		}
		_ = k8sClient.Delete(ctx, mcpGroup)
	})

	Context("when composite tools with parallel steps are configured", func() {
		It("should expose the composite tool in tool listing", func() {
			By("Creating and initializing MCP client for VirtualMCPServer")
			mcpClient, err := CreateInitializedMCPClient(vmcpNodePort, "toolhive-parallel-test", 30*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer mcpClient.Close()

			By("Listing tools from VirtualMCPServer")
			listRequest := mcp.ListToolsRequest{}
			tools, err := mcpClient.Client.ListTools(mcpClient.Ctx, listRequest)
			Expect(err).ToNot(HaveOccurred())

			By(fmt.Sprintf("VirtualMCPServer exposes %d tools", len(tools.Tools)))
			for _, tool := range tools.Tools {
				GinkgoWriter.Printf("  Tool: %s - %s\n", tool.Name, tool.Description)
			}

			// Should find the composite tool
			var foundComposite bool
			for _, tool := range tools.Tools {
				if tool.Name == compositeToolName {
					foundComposite = true
					Expect(tool.Description).To(ContainSubstring("parallel"))
					break
				}
			}
			Expect(foundComposite).To(BeTrue(), "Should find composite tool: %s", compositeToolName)

			// Should also have both backends' native echo tools (with prefix)
			foundBackends := make(map[string]bool)
			for _, tool := range tools.Tools {
				if tool.Name == fmt.Sprintf("%s_echo", backend1Name) {
					foundBackends[backend1Name] = true
				}
				if tool.Name == fmt.Sprintf("%s_echo", backend2Name) {
					foundBackends[backend2Name] = true
				}
			}
			Expect(foundBackends).To(HaveLen(2), "Should find both backend echo tools")
		})

		It("should execute parallel workflow and aggregate results", func() {
			By("Creating and initializing MCP client for VirtualMCPServer")
			mcpClient, err := CreateInitializedMCPClient(vmcpNodePort, "toolhive-parallel-test", 30*time.Second)
			Expect(err).ToNot(HaveOccurred())
			defer mcpClient.Close()

			By("Calling composite tool with test message")
			testMessage := "parallel_test_123"
			callRequest := mcp.CallToolRequest{}
			callRequest.Params.Name = compositeToolName
			callRequest.Params.Arguments = map[string]any{
				"message": testMessage,
			}

			result, err := mcpClient.Client.CallTool(mcpClient.Ctx, callRequest)
			Expect(err).ToNot(HaveOccurred(), "Composite tool call should succeed")
			Expect(result).ToNot(BeNil())
			Expect(result.Content).ToNot(BeEmpty(), "Should have content in response")

			// The result should contain combined outputs from both parallel steps
			// Final step combines: [backend1 result] + [backend2 result]
			GinkgoWriter.Printf("Parallel composite tool result: %+v\n", result.Content)
		})
	})

	Context("when verifying parallel workflow configuration", func() {
		It("should have correct composite tool spec with parallel steps", func() {
			vmcpServer := &mcpv1beta1.VirtualMCPServer{}
			err := k8sClient.Get(ctx, types.NamespacedName{
				Name:      vmcpServerName,
				Namespace: testNamespace,
			}, vmcpServer)
			Expect(err).ToNot(HaveOccurred())

			Expect(vmcpServer.Spec.Config.CompositeTools).To(HaveLen(1))

			compositeTool := vmcpServer.Spec.Config.CompositeTools[0]
			Expect(compositeTool.Name).To(Equal(compositeToolName))
			Expect(compositeTool.Steps).To(HaveLen(3))

			// Verify parallel steps (no dependencies)
			step1 := compositeTool.Steps[0]
			Expect(step1.ID).To(Equal("echo_backend1"))
			Expect(step1.DependsOn).To(BeEmpty(), "First step should have no dependencies (parallel)")

			step2 := compositeTool.Steps[1]
			Expect(step2.ID).To(Equal("echo_backend2"))
			Expect(step2.DependsOn).To(BeEmpty(), "Second step should have no dependencies (parallel)")

			// Verify final aggregation step depends on both parallel steps
			step3 := compositeTool.Steps[2]
			Expect(step3.ID).To(Equal("combine_results"))
			Expect(step3.DependsOn).To(ContainElements("echo_backend1", "echo_backend2"))

			// Verify template usage combines outputs from parallel steps
			step3Args := step3.Arguments.Value
			Expect(step3Args["input"]).To(ContainSubstring(".steps.echo_backend1"))
			Expect(step3Args["input"]).To(ContainSubstring(".steps.echo_backend2"))
		})

		It("should target different backends in parallel steps", func() {
			vmcpServer := &mcpv1beta1.VirtualMCPServer{}
			err := k8sClient.Get(ctx, types.NamespacedName{
				Name:      vmcpServerName,
				Namespace: testNamespace,
			}, vmcpServer)
			Expect(err).ToNot(HaveOccurred())

			compositeTool := vmcpServer.Spec.Config.CompositeTools[0]

			// Verify steps target different backends
			step1 := compositeTool.Steps[0]
			step2 := compositeTool.Steps[1]

			Expect(step1.Tool).To(ContainSubstring(backend1Name))
			Expect(step2.Tool).To(ContainSubstring(backend2Name))
		})
	})
})
