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

package e2e_test

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	"github.com/stacklok/toolhive/pkg/transport/types"
	"github.com/stacklok/toolhive/test/e2e"
)

func generateUniqueTelemetryServerName(prefix string) string {
	return fmt.Sprintf("%s-%d-%d-%d", prefix, os.Getpid(), time.Now().UnixNano(), GinkgoRandomSeed())
}

var _ = Describe("Telemetry Middleware E2E", Label("middleware", "telemetry", "e2e"), Serial, func() {
	var (
		config        *e2e.TestConfig
		proxyCmd      *exec.Cmd
		mcpServerName string
		workloadName  string
		transportType types.TransportType
		proxyMode     string
	)

	BeforeEach(func() {
		config = e2e.NewTestConfig()
		err := e2e.CheckTHVBinaryAvailable(config)
		Expect(err).ToNot(HaveOccurred())
		workloadName = generateUniqueTelemetryServerName("telemetry-test")
		mcpServerName = "osv" // Use OSV server as a reliable test server
		transportType = types.TransportTypeStreamableHTTP
	})

	JustBeforeEach(func() {
		// Build args for running the MCP server
		args := []string{"run", "--name", workloadName, "--transport", transportType.String()}

		if transportType == types.TransportTypeStdio {
			Expect(proxyMode).ToNot(BeEmpty())
			args = append(args, "--proxy-mode", proxyMode)
		}

		args = append(args, mcpServerName)

		By("Starting MCP server for telemetry testing")
		e2e.NewTHVCommand(config, args...).ExpectSuccess()

		err := e2e.WaitForMCPServer(config, workloadName, 60*time.Second)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		By("Cleaning up test resources")

		// Stop proxy if running
		if proxyCmd != nil && proxyCmd.Process != nil {
			proxyCmd.Process.Kill()
			proxyCmd.Wait()
		}

		// Stop and remove server
		if config.CleanupAfter {
			err := e2e.StopAndRemoveMCPServer(config, workloadName)
			Expect(err).ToNot(HaveOccurred(), "Should be able to stop and remove server")
		}
	})

	Context("when telemetry is enabled via environment variable", func() {
		BeforeEach(func() {
			// Enable telemetry via environment variable
			os.Setenv("TOOLHIVE_TELEMETRY_ENABLED", "true")
			os.Setenv("TOOLHIVE_TELEMETRY_SERVICE_NAME", "toolhive-e2e-test")
			os.Setenv("TOOLHIVE_TELEMETRY_SERVICE_VERSION", "test-1.0.0")
		})

		AfterEach(func() {
			// Clean up environment variables
			os.Unsetenv("TOOLHIVE_TELEMETRY_ENABLED")
			os.Unsetenv("TOOLHIVE_TELEMETRY_SERVICE_NAME")
			os.Unsetenv("TOOLHIVE_TELEMETRY_SERVICE_VERSION")
		})

		It("should capture telemetry data for MCP requests", func() {
			By("Starting the stdio proxy with telemetry enabled")
			stdin, outputBuffer := startProxyStdioForTelemetryTest(
				config,
				workloadName,
			)

			// Wait for proxy to start
			Eventually(func() string {
				return outputBuffer.String()
			}, 10*time.Second, 1*time.Second).Should(ContainSubstring("starting stdio proxy"))

			By("Sending MCP requests through the proxy")
			// Send an initialize request
			initRequest := map[string]interface{}{
				"jsonrpc": "2.0",
				"id":      "init-1",
				"method":  "initialize",
				"params": map[string]interface{}{
					"protocolVersion": "2024-11-05",
					"clientInfo": map[string]interface{}{
						"name":    "telemetry-test-client",
						"version": "1.0.0",
					},
				},
			}

			jsonRequest, err := json.Marshal(initRequest)
			Expect(err).ToNot(HaveOccurred())

			_, err = stdin.Write(jsonRequest)
			Expect(err).ToNot(HaveOccurred())
			_, err = stdin.Write([]byte("\n"))
			Expect(err).ToNot(HaveOccurred())

			// Send a tools/list request
			toolsListRequest := map[string]interface{}{
				"jsonrpc": "2.0",
				"id":      "tools-1",
				"method":  "tools/list",
			}

			jsonRequest, err = json.Marshal(toolsListRequest)
			Expect(err).ToNot(HaveOccurred())

			_, err = stdin.Write(jsonRequest)
			Expect(err).ToNot(HaveOccurred())
			_, err = stdin.Write([]byte("\n"))
			Expect(err).ToNot(HaveOccurred())

			// Wait a moment for telemetry to be processed
			time.Sleep(2 * time.Second)

			By("Verifying telemetry data was captured in logs")
			// Check that telemetry-related log entries exist
			logOutput := outputBuffer.String()

			// Look for telemetry indicators in the logs
			// The exact format may vary, but we should see some telemetry-related activity
			hasInitializeSpan := strings.Contains(logOutput, "initialize") ||
				strings.Contains(logOutput, "mcp.initialize") ||
				strings.Contains(logOutput, "span")

			hasToolsListSpan := strings.Contains(logOutput, "tools/list") ||
				strings.Contains(logOutput, "mcp.tools/list") ||
				strings.Contains(logOutput, "tools")

			// If telemetry is working, we should see some indication of spans or metrics
			hasTelemetryActivity := hasInitializeSpan || hasToolsListSpan ||
				strings.Contains(logOutput, "telemetry") ||
				strings.Contains(logOutput, "trace") ||
				strings.Contains(logOutput, "metric")

			if !hasTelemetryActivity {
				GinkgoWriter.Printf("Log output for telemetry verification:\n%s\n", logOutput)
			}

			// For now, just ensure the proxy worked correctly
			// The actual telemetry verification would require access to the telemetry backend
			Expect(strings.Contains(logOutput, "starting stdio proxy")).To(BeTrue())
		})

		It("should expose Prometheus metrics endpoint when enabled", func() {
			By("Enabling Prometheus metrics")
			os.Setenv("TOOLHIVE_TELEMETRY_PROMETHEUS_ENABLED", "true")
			os.Setenv("TOOLHIVE_TELEMETRY_PROMETHEUS_PORT", "9090")
			defer func() {
				os.Unsetenv("TOOLHIVE_TELEMETRY_PROMETHEUS_ENABLED")
				os.Unsetenv("TOOLHIVE_TELEMETRY_PROMETHEUS_PORT")
			}()

			By("Starting proxy with Prometheus metrics enabled")
			stdin, outputBuffer := startProxyStdioForTelemetryTest(
				config,
				workloadName,
			)

			// Wait for proxy to start
			Eventually(func() string {
				return outputBuffer.String()
			}, 10*time.Second, 1*time.Second).Should(ContainSubstring("starting stdio proxy"))

			By("Making MCP requests to generate metrics")
			// Send a simple tools/list request to generate metrics
			toolsRequest := map[string]interface{}{
				"jsonrpc": "2.0",
				"id":      "metrics-test",
				"method":  "tools/list",
			}

			jsonRequest, err := json.Marshal(toolsRequest)
			Expect(err).ToNot(HaveOccurred())

			_, err = stdin.Write(jsonRequest)
			Expect(err).ToNot(HaveOccurred())
			_, err = stdin.Write([]byte("\n"))
			Expect(err).ToNot(HaveOccurred())

			// Wait for metrics to be recorded
			time.Sleep(3 * time.Second)

			By("Attempting to access Prometheus metrics endpoint")
			// Try to access the metrics endpoint
			// Note: This is a best-effort test since the exact port binding might vary
			possiblePorts := []string{"9090", "8080", "8081", "9091"}
			metricsFound := false

			for _, port := range possiblePorts {
				metricsURL := fmt.Sprintf("http://localhost:%s/metrics", port)
				resp, err := http.Get(metricsURL)
				if err != nil {
					continue
				}
				defer resp.Body.Close()

				if resp.StatusCode == http.StatusOK {
					body, err := io.ReadAll(resp.Body)
					if err == nil && len(body) > 0 {
						metricsContent := string(body)
						GinkgoWriter.Printf("Found metrics on port %s:\n%s\n", port, metricsContent)

						// Look for ToolHive-specific metrics
						if strings.Contains(metricsContent, "toolhive_mcp") {
							metricsFound = true
							break
						}
					}
				}
			}

			// For now, we'll just verify the proxy is working
			// The actual metrics endpoint testing would require more specific setup
			logOutput := outputBuffer.String()
			Expect(strings.Contains(logOutput, "starting stdio proxy")).To(BeTrue())

			if metricsFound {
				GinkgoWriter.Println("Successfully found ToolHive metrics endpoint")
			} else {
				GinkgoWriter.Println("Metrics endpoint not accessible (this may be expected in test environment)")
			}
		})
	})

	Context("when telemetry environment variables are set", func() {
		It("should respect custom environment variable configurations", func() {
			By("Setting custom environment variables for telemetry")
			os.Setenv("CUSTOM_ENV_VAR", "test-value")
			os.Setenv("TOOLHIVE_TELEMETRY_ENVIRONMENT_VARIABLES", "CUSTOM_ENV_VAR,USER")
			defer func() {
				os.Unsetenv("CUSTOM_ENV_VAR")
				os.Unsetenv("TOOLHIVE_TELEMETRY_ENVIRONMENT_VARIABLES")
			}()

			By("Starting proxy with environment variable telemetry")
			stdin, outputBuffer := startProxyStdioForTelemetryTest(
				config,
				workloadName,
			)

			// Wait for proxy to start
			Eventually(func() string {
				return outputBuffer.String()
			}, 20*time.Second, 1*time.Second).Should(ContainSubstring("starting stdio proxy"))

			By("Sending MCP request")
			request := map[string]interface{}{
				"jsonrpc": "2.0",
				"id":      "env-test",
				"method":  "tools/list",
			}

			jsonRequest, err := json.Marshal(request)
			Expect(err).ToNot(HaveOccurred())

			_, err = stdin.Write(jsonRequest)
			Expect(err).ToNot(HaveOccurred())
			_, err = stdin.Write([]byte("\n"))
			Expect(err).ToNot(HaveOccurred())

			// Wait for processing
			time.Sleep(2 * time.Second)

			// Verify the proxy worked
			logOutput := outputBuffer.String()
			Expect(strings.Contains(logOutput, "starting stdio proxy")).To(BeTrue())
		})
	})
})

// startProxyStdioForTelemetryTest starts a stdio proxy for telemetry testing
// and returns the command, stdin pipe, and output buffer for monitoring
func startProxyStdioForTelemetryTest(config *e2e.TestConfig, workloadName string) (io.WriteCloser, *SafeBuffer) {
	By("starting stdio proxy with telemetry")

	// Get the server URL first
	serverURL, err := e2e.GetMCPServerURL(config, workloadName)
	Expect(err).ToNot(HaveOccurred())

	// Extract base URL for transparent proxy
	// With streamable-http: http://127.0.0.1:PORT/mcp (no fragment)
	baseURL := strings.TrimSuffix(serverURL, "/mcp")
	GinkgoWriter.Printf("Base URL for telemetry proxy: %s\n", baseURL)

	// Start the proxy command
	cmd := exec.Command(config.THVBinary, "proxy", "stdio", workloadName, "--debug") //nolint:gosec
	cmd.Env = os.Environ()

	// Create pipes for stdin and stdout/stderr
	stdin, err := cmd.StdinPipe()
	Expect(err).ToNot(HaveOccurred())

	stdout, err := cmd.StdoutPipe()
	Expect(err).ToNot(HaveOccurred())

	stderr, err := cmd.StderrPipe()
	Expect(err).ToNot(HaveOccurred())

	// Start the command
	err = cmd.Start()
	Expect(err).ToNot(HaveOccurred())

	// Create a buffer to capture output
	outputBuffer := NewSafeBuffer()

	// Start goroutines to capture output
	go func() {
		defer GinkgoRecover()
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			line := scanner.Text()
			GinkgoWriter.Printf("STDOUT: %s\n", line)
			outputBuffer.WriteString(line + "\n")
		}
	}()

	go func() {
		defer GinkgoRecover()
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			line := scanner.Text()
			GinkgoWriter.Printf("STDERR: %s\n", line)
			outputBuffer.WriteString(line + "\n")
		}
	}()

	return stdin, outputBuffer
}

// SafeBuffer is a thread-safe string buffer for capturing output
type SafeBuffer struct {
	buffer strings.Builder
}

func NewSafeBuffer() *SafeBuffer {
	return &SafeBuffer{}
}

func (sb *SafeBuffer) WriteString(s string) {
	sb.buffer.WriteString(s)
}

func (sb *SafeBuffer) String() string {
	return sb.buffer.String()
}

// Additional helper functions for telemetry verification
