package mcpcheck

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"runtime"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
)

func TestCheckHTTPMCPServerListsTools(t *testing.T) {
	initialized := false
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "Bearer test", r.Header.Get("Authorization"))
		var req struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
		}
		require.NoError(t, json.NewDecoder(r.Body).Decode(&req))
		switch req.Method {
		case "initialize":
			w.Header().Set("Mcp-Session-Id", "test-session")
			_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"test","version":"1"}}}`))
		case "notifications/initialized":
			require.Equal(t, "test-session", r.Header.Get("Mcp-Session-Id"))
			initialized = true
			w.WriteHeader(http.StatusAccepted)
		case "tools/list":
			require.True(t, initialized, "tools/list must be sent after notifications/initialized")
			require.Equal(t, "test-session", r.Header.Get("Mcp-Session-Id"))
			_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":2,"result":{"tools":[{"name":"hello","description":"Hello tool"}]}}`))
		default:
			w.WriteHeader(http.StatusBadRequest)
		}
	}))
	defer server.Close()

	check := Check(context.Background(), profile.MCPServer{
		ID:          "http_one",
		DisplayName: "HTTP One",
		Enabled:     true,
		Transport:   profile.MCPServerTransportHTTP,
		HTTP:        &profile.MCPHTTPConfig{URL: server.URL, Headers: []string{"Authorization: Bearer test"}},
	})

	require.True(t, check.OK, check.Message)
	require.Equal(t, "tools", check.Stage)
	require.Len(t, check.Tools, 1)
	require.Equal(t, "hello", check.Tools[0].Name)
}

func TestCheckHTTPMCPServerResolvesEnvVar(t *testing.T) {
	t.Setenv("MCP_CHECK_TOKEN", "Bearer env-resolved")
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "Bearer env-resolved", r.Header.Get("Authorization"))
		var req struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
		}
		require.NoError(t, json.NewDecoder(r.Body).Decode(&req))
		switch req.Method {
		case "initialize":
			w.Header().Set("Mcp-Session-Id", "test-session")
			_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"test","version":"1"}}}`))
		case "notifications/initialized":
			w.WriteHeader(http.StatusAccepted)
		case "tools/list":
			_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":2,"result":{"tools":[]}}`))
		default:
			w.WriteHeader(http.StatusBadRequest)
		}
	}))
	defer server.Close()

	check := Check(context.Background(), profile.MCPServer{
		ID:          "http_env",
		DisplayName: "HTTP Env",
		Enabled:     true,
		Transport:   profile.MCPServerTransportHTTP,
		HTTP:        &profile.MCPHTTPConfig{URL: server.URL, Headers: []string{"Authorization: $MCP_CHECK_TOKEN"}},
	})

	require.True(t, check.OK, check.Message)
}

func TestResolveHTTPHeadersResolvesShellSnapshotEnvVar(t *testing.T) {
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{
		"MCP_CHECK_SHELL_TOKEN": "Bearer shell-env-resolved",
	})
	t.Cleanup(func() {
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})

	headers, err := resolveHTTPHeaders([]string{"Authorization: $MCP_CHECK_SHELL_TOKEN"})
	require.NoError(t, err)
	require.Equal(t, map[string]string{"Authorization": "Bearer shell-env-resolved"}, headers)
}

func TestCheckHTTPMCPServerClassifiesAuthStatus(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = w.Write([]byte("missing token"))
	}))
	defer server.Close()

	check := Check(context.Background(), profile.MCPServer{
		ID:          "http_one",
		DisplayName: "HTTP One",
		Enabled:     true,
		Transport:   profile.MCPServerTransportHTTP,
		HTTP:        &profile.MCPHTTPConfig{URL: server.URL},
	})

	require.False(t, check.OK)
	require.Equal(t, "auth", check.Stage)
	require.Contains(t, check.Message, "401")
}

func TestCheckStdioMCPServerListsTools(t *testing.T) {
	dir := t.TempDir()
	script := filepath.Join(dir, "server.sh")
	require.NoError(t, os.WriteFile(script, []byte(`#!/bin/sh
count=0
while IFS= read -r line; do
  count=$((count + 1))
  if [ "$count" = "1" ]; then
    printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"stdio","version":"1"}}}'
  elif [ "$count" = "3" ]; then
    printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{"tools":[{"name":"stdio_tool","description":"Stdio tool"}]}}'
    exit 0
  fi
done
`), 0o755))

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	check := Check(ctx, profile.MCPServer{
		ID:          "stdio_one",
		DisplayName: "Stdio One",
		Enabled:     true,
		Transport:   profile.MCPServerTransportStdio,
		Stdio:       &profile.MCPStdioConfig{Command: "/bin/sh", Args: []string{script}, CWD: dir},
	})

	require.True(t, check.OK, check.Message)
	require.Equal(t, "stdio_tool", check.Tools[0].Name)
}

func TestCheckStdioMCPServerUsesShellPathAndEnv(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("test uses a POSIX shell script")
	}

	dir := t.TempDir()
	script := filepath.Join(dir, "mcp-shell-server")
	require.NoError(t, os.WriteFile(script, []byte(`#!/bin/sh
if [ "$MCP_SECRET" != "secret-from-shell" ]; then
  echo "missing MCP_SECRET" >&2
  exit 2
fi
count=0
while IFS= read -r line; do
  count=$((count + 1))
  if [ "$count" = "1" ]; then
    printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"stdio","version":"1"}}}'
  elif [ "$count" = "3" ]; then
    printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{"tools":[{"name":"shell_tool"}]}}'
    exit 0
  fi
done
`), 0o755))

	runtimeconfig.SetShellPathForRuntimeResolution(dir)
	runtimeconfig.SetShellEnvForRuntimeResolution(map[string]string{
		"MCP_SECRET": "secret-from-shell",
	})
	t.Cleanup(func() {
		runtimeconfig.SetShellPathForRuntimeResolution("")
		runtimeconfig.SetShellEnvForRuntimeResolution(nil)
	})
	t.Setenv("PATH", t.TempDir())

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	check := Check(ctx, profile.MCPServer{
		ID:          "stdio_shell",
		DisplayName: "Stdio Shell",
		Enabled:     true,
		Transport:   profile.MCPServerTransportStdio,
		Stdio: &profile.MCPStdioConfig{
			Command: "mcp-shell-server",
			Env:     []string{"MCP_SECRET=$MCP_SECRET"},
		},
	})

	require.True(t, check.OK, check.Message)
	require.Equal(t, "shell_tool", check.Tools[0].Name)
}