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

// TODO: Follow-up issue tracks infra-heavy tests that need JWT/OIDC, circuit
// breaker, Redis session, and Tier-2 TEI embedding infrastructure:
// https://github.com/stacklok/toolhive/issues/4944

package e2e_test

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	"gopkg.in/yaml.v3"

	thvjson "github.com/stacklok/toolhive/pkg/json"
	"github.com/stacklok/toolhive/pkg/vmcp"
	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
	"github.com/stacklok/toolhive/test/e2e"
)

// modifyVMCPConfig reads a vMCP YAML config file, calls fn to mutate it, then
// writes it back. Used to layer feature-specific settings on top of a base
// config generated by `thv vmcp init`.
func modifyVMCPConfig(path string, fn func(*vmcpconfig.Config)) error {
	data, err := os.ReadFile(path) //nolint:gosec // path is test-controlled
	if err != nil {
		return fmt.Errorf("reading vmcp config %s: %w", path, err)
	}
	var cfg vmcpconfig.Config
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return fmt.Errorf("unmarshalling vmcp config: %w", err)
	}
	fn(&cfg)
	out, err := yaml.Marshal(&cfg)
	if err != nil {
		return fmt.Errorf("marshalling vmcp config: %w", err)
	}
	if err := os.WriteFile(path, out, 0o600); err != nil {
		return fmt.Errorf("writing vmcp config %s: %w", path, err)
	}
	return nil
}

// singleBackendFixture holds shared state for feature contexts that use one
// yardstick backend. Call setup in BeforeEach and teardown in AfterEach.
type singleBackendFixture struct {
	cfg         *e2e.TestConfig
	groupName   string
	backendName string
	vMCPCmd     *exec.Cmd
	vMCPPort    int
	tmpDir      string // empty when withTmpDir is false
}

func (f *singleBackendFixture) setup(groupPrefix, tmpDirPattern string) {
	f.cfg = e2e.NewTestConfig()
	f.groupName = e2e.GenerateUniqueServerName(groupPrefix)
	f.backendName = e2e.GenerateUniqueServerName("yardstick")
	f.vMCPPort = allocateVMCPPort()

	if tmpDirPattern != "" {
		var err error
		f.tmpDir, err = os.MkdirTemp("", tmpDirPattern)
		Expect(err).ToNot(HaveOccurred())
		DeferCleanup(func() { _ = os.RemoveAll(f.tmpDir) })
	}

	Expect(e2e.CheckTHVBinaryAvailable(f.cfg)).To(Succeed())

	e2e.NewTHVCommand(f.cfg, "group", "create", f.groupName).ExpectSuccess()
	startYardstick(f.cfg, f.groupName, f.backendName)
}

func (f *singleBackendFixture) teardown() {
	if f.cfg == nil {
		return
	}
	stopVMCPProcess(f.vMCPCmd)
	if f.cfg.CleanupAfter {
		if err := e2e.StopAndRemoveMCPServer(f.cfg, f.backendName); err != nil {
			GinkgoWriter.Printf("cleanup: StopAndRemoveMCPServer(%s) failed: %v\n", f.backendName, err)
		}
		if err := e2e.RemoveGroup(f.cfg, f.groupName); err != nil {
			GinkgoWriter.Printf("cleanup: RemoveGroup(%s) failed: %v\n", f.groupName, err)
		}
	}
}

// twoBackendFixture holds shared state for feature contexts that use two
// yardstick backends.
//
// Backend lifecycle (group + containers) is separated from per-test state
// (vMCP port, process, tmpDir) so that multiple test contexts can share the
// same running backends via BeforeAll/AfterAll while each test still gets its
// own vMCP serve process.
type twoBackendFixture struct {
	cfg          *e2e.TestConfig
	groupName    string
	backendAName string
	backendBName string
	vMCPCmd      *exec.Cmd
	vMCPPort     int
	tmpDir       string
}

// setupBackends starts the group and both yardstick containers. Call once in
// BeforeAll when sharing backends across multiple tests.
func (f *twoBackendFixture) setupBackends(groupPrefix string) {
	f.cfg = e2e.NewTestConfig()
	f.groupName = e2e.GenerateUniqueServerName(groupPrefix)
	f.backendAName = e2e.GenerateUniqueServerName("yardstick-a")
	f.backendBName = e2e.GenerateUniqueServerName("yardstick-b")

	Expect(e2e.CheckTHVBinaryAvailable(f.cfg)).To(Succeed())

	By("creating group and two yardstick backends")
	e2e.NewTHVCommand(f.cfg, "group", "create", f.groupName).ExpectSuccess()
	// Use different ports so both containers can bind successfully.
	// yardstick does not honour MCP_PORT; its listening port must be set
	// explicitly via the -port flag and matched by --target-port.
	startYardstickOnPort(f.cfg, f.groupName, f.backendAName, 8080)
	startYardstickOnPort(f.cfg, f.groupName, f.backendBName, 8081)
}

// teardownBackends stops and removes the group and both backends. Call in
// AfterAll to match setupBackends.
func (f *twoBackendFixture) teardownBackends() {
	if f.cfg == nil {
		return
	}
	if f.cfg.CleanupAfter {
		if err := e2e.StopAndRemoveMCPServer(f.cfg, f.backendAName); err != nil {
			GinkgoWriter.Printf("cleanup: StopAndRemoveMCPServer(%s) failed: %v\n", f.backendAName, err)
		}
		if err := e2e.StopAndRemoveMCPServer(f.cfg, f.backendBName); err != nil {
			GinkgoWriter.Printf("cleanup: StopAndRemoveMCPServer(%s) failed: %v\n", f.backendBName, err)
		}
		if err := e2e.RemoveGroup(f.cfg, f.groupName); err != nil {
			GinkgoWriter.Printf("cleanup: RemoveGroup(%s) failed: %v\n", f.groupName, err)
		}
	}
}

// setupPerTest allocates a fresh port and tmpDir for one test. Call in
// BeforeEach when backends are already running.
func (f *twoBackendFixture) setupPerTest(tmpDirPattern string) {
	f.vMCPPort = allocateVMCPPort()
	var err error
	f.tmpDir, err = os.MkdirTemp("", tmpDirPattern)
	Expect(err).ToNot(HaveOccurred())
	DeferCleanup(func() { _ = os.RemoveAll(f.tmpDir) })
}

// teardownPerTest stops the vMCP serve process. Call in AfterEach.
func (f *twoBackendFixture) teardownPerTest() {
	stopVMCPProcess(f.vMCPCmd)
	f.vMCPCmd = nil
}

var _ = Describe("vMCP CLI features", Label("vmcp", "e2e", "features"), func() {

	// -------------------------------------------------------------------------
	// Contexts 1 & 2: Two-backend scenarios.
	// Both tests share the same group and yardstick containers (started once
	// in BeforeAll) to avoid paying the container-startup cost twice.
	// Each test still gets its own vMCP port, process, and tmpDir.
	// -------------------------------------------------------------------------
	Describe("two-backend scenarios", Ordered, func() {
		var fx twoBackendFixture

		BeforeAll(func() { fx.setupBackends("vmcp-feat-shared") })
		AfterAll(func() { fx.teardownBackends() })

		// Context 1: Conflict resolution — prefix strategy
		// Two identical yardstick backends expose the same tool names. With prefix
		// strategy enabled, both sets of tools appear with backend-name prefixes.
		Context("conflict resolution (prefix strategy)", func() {
			BeforeEach(func() { fx.setupPerTest("vmcp-feat-conflict-*") })
			AfterEach(func() { fx.teardownPerTest() })

			It("prefixes tool names with backend names so both echo tools are visible", func() {
				configPath := filepath.Join(fx.tmpDir, "vmcp.yaml")
				initVMCPConfig(fx.cfg, fx.groupName, configPath)

				Expect(modifyVMCPConfig(configPath, func(c *vmcpconfig.Config) {
					if c.Aggregation == nil {
						c.Aggregation = &vmcpconfig.AggregationConfig{}
					}
					c.Aggregation.ConflictResolution = vmcp.ConflictStrategyPrefix
					c.Aggregation.ConflictResolutionConfig = &vmcpconfig.ConflictResolutionConfig{
						PrefixFormat: "{workload}_",
					}
				})).To(Succeed())

				By("starting vMCP serve with prefix config")
				fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg,
					"vmcp", "serve",
					"--config", configPath,
					"--port", fmt.Sprintf("%d", fx.vMCPPort),
				)
				vMCPURL := fmt.Sprintf("http://127.0.0.1:%d/mcp", fx.vMCPPort)
				Expect(e2e.WaitForMCPServerReady(fx.cfg, vMCPURL, "streamable-http", 60*time.Second)).To(Succeed())

				By("listing tools and verifying prefix strategy")
				ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
				defer cancel()
				mcpClient, err := e2e.NewMCPClientForStreamableHTTP(fx.cfg, vMCPURL)
				Expect(err).ToNot(HaveOccurred())
				defer func() { _ = mcpClient.Close() }()
				Expect(mcpClient.Initialize(ctx)).To(Succeed())

				tools, err := mcpClient.ListTools(ctx)
				Expect(err).ToNot(HaveOccurred())

				var echoTools []string
				for _, t := range tools.Tools {
					if strings.Contains(t.Name, "echo") {
						echoTools = append(echoTools, t.Name)
					}
				}
				By(fmt.Sprintf("found echo tools: %v", echoTools))
				Expect(echoTools).To(HaveLen(2), "prefix strategy should expose one echo tool per backend")

				hasA := false
				hasB := false
				for _, name := range echoTools {
					if strings.HasPrefix(name, fx.backendAName+"_") {
						hasA = true
					}
					if strings.HasPrefix(name, fx.backendBName+"_") {
						hasB = true
					}
				}
				Expect(hasA).To(BeTrue(), fmt.Sprintf("expected echo tool prefixed with %s_", fx.backendAName))
				Expect(hasB).To(BeTrue(), fmt.Sprintf("expected echo tool prefixed with %s_", fx.backendBName))
			})
		})

		// Context 2: Tool filtering — per-workload allow-list + ExcludeAll
		// backendA exposes only "echo"; backendB is completely hidden.
		Context("tool filtering (per-workload filter + ExcludeAll)", func() {
			BeforeEach(func() { fx.setupPerTest("vmcp-feat-filter-*") })
			AfterEach(func() { fx.teardownPerTest() })

			It("only exposes the filtered tool from backendA; backendB tools are hidden", func() {
				configPath := filepath.Join(fx.tmpDir, "vmcp.yaml")
				initVMCPConfig(fx.cfg, fx.groupName, configPath)

				Expect(modifyVMCPConfig(configPath, func(c *vmcpconfig.Config) {
					if c.Aggregation == nil {
						c.Aggregation = &vmcpconfig.AggregationConfig{}
					}
					// ConflictResolution is required by the aggregation validator even when
					// per-workload filters are the focus of the test.
					c.Aggregation.ConflictResolution = vmcp.ConflictStrategyPrefix
					c.Aggregation.ConflictResolutionConfig = &vmcpconfig.ConflictResolutionConfig{
						PrefixFormat: "{workload}_",
					}
					c.Aggregation.Tools = []*vmcpconfig.WorkloadToolConfig{
						{Workload: fx.backendAName, Filter: []string{"echo"}},
						{Workload: fx.backendBName, ExcludeAll: true},
					}
				})).To(Succeed())

				By("starting vMCP serve with filter config")
				fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg,
					"vmcp", "serve",
					"--config", configPath,
					"--port", fmt.Sprintf("%d", fx.vMCPPort),
				)
				vMCPURL := fmt.Sprintf("http://127.0.0.1:%d/mcp", fx.vMCPPort)
				Expect(e2e.WaitForMCPServerReady(fx.cfg, vMCPURL, "streamable-http", 60*time.Second)).To(Succeed())

				ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
				defer cancel()
				mcpClient, err := e2e.NewMCPClientForStreamableHTTP(fx.cfg, vMCPURL)
				Expect(err).ToNot(HaveOccurred())
				defer func() { _ = mcpClient.Close() }()
				Expect(mcpClient.Initialize(ctx)).To(Succeed())

				tools, err := mcpClient.ListTools(ctx)
				Expect(err).ToNot(HaveOccurred())

				By(fmt.Sprintf("verifying tool visibility; got %d tools", len(tools.Tools)))
				var echoCount int
				for _, t := range tools.Tools {
					Expect(t.Name).ToNot(ContainSubstring(fx.backendBName),
						"backendB tools must be hidden via ExcludeAll")
					if strings.Contains(t.Name, "echo") {
						echoCount++
					}
				}
				Expect(echoCount).To(Equal(1), "only backendA echo should be visible")
			})
		})
	})

	// -------------------------------------------------------------------------
	// Context 3: Global ExcludeAll — all backend tools hidden
	// tools/list must succeed but return an empty slice.
	// -------------------------------------------------------------------------
	Context("global ExcludeAll (all tools hidden)", func() {
		var fx singleBackendFixture

		BeforeEach(func() { fx.setup("vmcp-feat-excludeall", "vmcp-feat-excludeall-*") })
		AfterEach(func() { fx.teardown() })

		It("returns an empty tools list when ExcludeAllTools is true", func() {
			configPath := filepath.Join(fx.tmpDir, "vmcp.yaml")
			initVMCPConfig(fx.cfg, fx.groupName, configPath)

			Expect(modifyVMCPConfig(configPath, func(c *vmcpconfig.Config) {
				if c.Aggregation == nil {
					c.Aggregation = &vmcpconfig.AggregationConfig{}
				}
				c.Aggregation.ConflictResolution = vmcp.ConflictStrategyPrefix
				c.Aggregation.ExcludeAllTools = true
			})).To(Succeed())

			By("starting vMCP serve with ExcludeAllTools")
			fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg,
				"vmcp", "serve",
				"--config", configPath,
				"--port", fmt.Sprintf("%d", fx.vMCPPort),
			)
			vMCPURL := fmt.Sprintf("http://127.0.0.1:%d/mcp", fx.vMCPPort)
			Expect(e2e.WaitForMCPServerReady(fx.cfg, vMCPURL, "streamable-http", 60*time.Second)).To(Succeed())

			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()
			mcpClient, err := e2e.NewMCPClientForStreamableHTTP(fx.cfg, vMCPURL)
			Expect(err).ToNot(HaveOccurred())
			defer func() { _ = mcpClient.Close() }()
			Expect(mcpClient.Initialize(ctx)).To(Succeed())

			tools, err := mcpClient.ListTools(ctx)
			Expect(err).ToNot(HaveOccurred(), "tools/list must succeed even when all tools are hidden")
			Expect(tools.Tools).To(BeEmpty(), "ExcludeAllTools should hide all backend tools")
		})
	})

	// -------------------------------------------------------------------------
	// Context 4: Composite sequential tool
	// An "echo_twice" composite calls yardstick's echo tool twice in sequence.
	// Verifies that: (a) the composite appears in tools/list, (b) calling it
	// with a message argument succeeds and returns non-empty content.
	// -------------------------------------------------------------------------
	Context("composite sequential tool", func() {
		var fx singleBackendFixture

		BeforeEach(func() { fx.setup("vmcp-feat-composite", "vmcp-feat-composite-*") })
		AfterEach(func() { fx.teardown() })

		It("exposes the composite tool in tools/list and executes it successfully", func() {
			configPath := filepath.Join(fx.tmpDir, "vmcp.yaml")
			initVMCPConfig(fx.cfg, fx.groupName, configPath)

			Expect(modifyVMCPConfig(configPath, func(c *vmcpconfig.Config) {
				if c.Aggregation == nil {
					c.Aggregation = &vmcpconfig.AggregationConfig{}
				}
				// ConflictResolution is required by the aggregation validator.
				c.Aggregation.ConflictResolution = vmcp.ConflictStrategyPrefix

				c.CompositeTools = []vmcpconfig.CompositeToolConfig{
					{
						Name:        "echo_twice",
						Description: "Echoes the input message twice in sequence",
						Parameters: thvjson.NewMap(map[string]any{
							"type": "object",
							"properties": map[string]any{
								"message": map[string]any{
									"type":        "string",
									"description": "The message to echo twice",
								},
							},
							"required": []any{"message"},
						}),
						Steps: []vmcpconfig.WorkflowStepConfig{
							{
								ID:   "first_echo",
								Type: "tool",
								Tool: fmt.Sprintf("%s.echo", fx.backendName),
								Arguments: thvjson.NewMap(map[string]any{
									"input": "{{ .params.message }}",
								}),
							},
							{
								ID:        "second_echo",
								Type:      "tool",
								Tool:      fmt.Sprintf("%s.echo", fx.backendName),
								DependsOn: []string{"first_echo"},
								Arguments: thvjson.NewMap(map[string]any{
									"input": "{{ .params.message }}",
								}),
							},
						},
					},
				}
			})).To(Succeed())

			By("starting vMCP serve with composite tool config")
			fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg,
				"vmcp", "serve",
				"--config", configPath,
				"--port", fmt.Sprintf("%d", fx.vMCPPort),
			)
			vMCPURL := fmt.Sprintf("http://127.0.0.1:%d/mcp", fx.vMCPPort)
			Expect(e2e.WaitForMCPServerReady(fx.cfg, vMCPURL, "streamable-http", 60*time.Second)).To(Succeed())

			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()
			mcpClient, err := e2e.NewMCPClientForStreamableHTTP(fx.cfg, vMCPURL)
			Expect(err).ToNot(HaveOccurred())
			defer func() { _ = mcpClient.Close() }()
			Expect(mcpClient.Initialize(ctx)).To(Succeed())

			By("verifying echo_twice appears in tools/list")
			tools, err := mcpClient.ListTools(ctx)
			Expect(err).ToNot(HaveOccurred())

			var foundComposite bool
			for _, t := range tools.Tools {
				if t.Name == "echo_twice" {
					foundComposite = true
					Expect(t.Description).To(Equal("Echoes the input message twice in sequence"))
					break
				}
			}
			Expect(foundComposite).To(BeTrue(), "echo_twice composite tool must appear in tools/list")

			By("calling echo_twice and verifying non-empty result")
			result, err := mcpClient.CallTool(ctx, "echo_twice", map[string]any{
				"message": "hellocomposite",
			})
			Expect(err).ToNot(HaveOccurred(), "composite tool call must succeed")
			Expect(result.IsError).To(BeFalse(), "composite tool must not return an error result")
			Expect(result.Content).ToNot(BeEmpty(), "composite tool must return content")
			Expect(mcp.GetTextFromContent(result.Content[0])).To(ContainSubstring("hellocomposite"),
				"composite tool result must contain the echoed message")
		})
	})

	// -------------------------------------------------------------------------
	// Context 5: Tier-1 optimizer (FTS5) — quick mode
	// `thv vmcp serve --group <name> --optimizer` must expose exactly two tools:
	// find_tool and call_tool. Calling find_tool with a query must return results.
	// -------------------------------------------------------------------------
	Context("Tier-1 optimizer (--optimizer flag, quick mode)", func() {
		var fx singleBackendFixture

		BeforeEach(func() { fx.setup("vmcp-feat-optimizer", "") })
		AfterEach(func() { fx.teardown() })

		It("exposes only find_tool and call_tool when --optimizer is set", func() {
			By("starting vMCP serve in quick mode with --optimizer")
			fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg,
				"vmcp", "serve",
				"--group", fx.groupName,
				"--optimizer",
				"--port", fmt.Sprintf("%d", fx.vMCPPort),
			)
			vMCPURL := fmt.Sprintf("http://127.0.0.1:%d/mcp", fx.vMCPPort)
			Expect(e2e.WaitForMCPServerReady(fx.cfg, vMCPURL, "streamable-http", 60*time.Second)).To(Succeed())

			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()
			mcpClient, err := e2e.NewMCPClientForStreamableHTTP(fx.cfg, vMCPURL)
			Expect(err).ToNot(HaveOccurred())
			defer func() { _ = mcpClient.Close() }()
			Expect(mcpClient.Initialize(ctx)).To(Succeed())

			By("verifying only find_tool and call_tool are exposed")
			tools, err := mcpClient.ListTools(ctx)
			Expect(err).ToNot(HaveOccurred())
			Expect(tools.Tools).To(HaveLen(2), "optimizer mode must expose exactly 2 tools")

			names := make([]string, len(tools.Tools))
			for i, t := range tools.Tools {
				names[i] = t.Name
			}
			Expect(names).To(ConsistOf("find_tool", "call_tool"))

			By("calling find_tool to verify it returns results")
			result, err := mcpClient.CallTool(ctx, "find_tool", map[string]any{
				"tool_description": "echo a message",
			})
			Expect(err).ToNot(HaveOccurred(), "find_tool must succeed")
			Expect(result.IsError).To(BeFalse(), "find_tool must not return an error result")
			Expect(result.Content).ToNot(BeEmpty(), "find_tool must return tool suggestions")
		})
	})

}) // end Describe("vMCP CLI features")
