package backend

import (
	"context"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/client/transport"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
)

func TestSetupProxy(t *testing.T) {
	logger := zap.NewNop()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Create a test MCP server
	s := server.NewMCPServer("test", "1.0")
	httpServer := server.NewStreamableHTTPServer(
		s,
		server.WithEndpointPath("/mcp"),
		server.WithStateLess(true),
	)

	// Start HTTP test server
	testServer := httptest.NewServer(httpServer)
	defer testServer.Close()

	// Create HTTP transport client
	serverAddr := testServer.URL
	mcpEndpoint := serverAddr + "/mcp"

	tr, err := transport.NewStreamableHTTP(mcpEndpoint)
	require.NoError(t, err, "failed to create HTTP transport")

	// Test the setupProxy function with HTTP transport
	c, handler, err := setupProxy(ctx, logger, tr)
	require.NoError(t, err, "setupProxy should not return error")
	require.NotNil(t, c, "client should not be nil")
	require.NotNil(t, handler, "handler should not be nil")

	// Verify the client can actually connect
	clientConn := client.NewClient(tr)
	defer clientConn.Close()

	err = clientConn.Start(ctx)
	require.NoError(t, err, "client should start successfully")

	init, err := clientConn.Initialize(ctx, mcp.InitializeRequest{})
	require.NoError(t, err, "client should initialize successfully")
	require.Equal(t, "test", init.ServerInfo.Name)
	require.Equal(t, "1.0", init.ServerInfo.Version)
}

func TestProxyBackendRunWithInvalidCommand(t *testing.T) {
	logger := zap.NewNop()
	pb := NewProxyBackend(logger, []string{"sh", "-c", "exit 0"})
	defer pb.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	handler, err := pb.Run(ctx)
	require.Error(t, err, "expected error for invalid command")
	require.Nil(t, handler, "handler should be nil on error")
}

func TestProxyBackendRun(t *testing.T) {
	logger := zap.NewNop()
	cmd := []string{"go", "run", "./testserver"}
	pb := NewProxyBackend(logger, cmd)
	defer pb.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)

	handler, err := pb.Run(ctx)
	require.NoError(t, err, "Run should not return error")
	require.NotNil(t, handler, "handler should not be nil")

	checkCh := make(chan struct{})
	go func() {
		pb.Wait()
		close(checkCh)
	}()

	timeout := time.After(10 * time.Millisecond)
	select {
	case <-checkCh:
		t.Error("Test completed too early")
	case <-timeout:
		// Test timed out
	}

	cancel()

	timeout = time.After(10 * time.Second)
	select {
	case <-checkCh:
		// Test completed successfully
	case <-timeout:
		t.Error("Test timed out")
	}
}