package mcp_test

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner/mcp"
)

func TestClientConcurrentInvocationsIsolateCredentialsSessionsAndBudgets(
	t *testing.T,
) {
	t.Parallel()

	var (
		serverErrorsMu sync.Mutex
		serverErrors   []error
	)
	recordServerError := func(format string, arguments ...any) {
		serverErrorsMu.Lock()
		defer serverErrorsMu.Unlock()
		serverErrors = append(
			serverErrors,
			fmt.Errorf(format, arguments...),
		)
	}
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		credential := request.Header.Get("X-MCP-Key")
		if credential == "" {
			recordServerError("request %s had no credential", request.Method)
			writer.WriteHeader(http.StatusUnauthorized)
			return
		}
		expectedSession := "session-" + credential
		if request.Method == http.MethodDelete {
			if got := request.Header.Get("Mcp-Session-Id"); got != expectedSession {
				recordServerError(
					"credential %q DELETE session = %q, want %q",
					credential,
					got,
					expectedSession,
				)
			}
			writer.WriteHeader(http.StatusNoContent)
			return
		}

		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			if got := request.Header.Get("Mcp-Session-Id"); got != "" {
				recordServerError(
					"credential %q discover carried session %q",
					credential,
					got,
				)
			}
			writeRPCError(t, writer, rpcRequest.ID, -32601, "legacy")
		case "initialize":
			if got := request.Header.Get("Mcp-Session-Id"); got != "" {
				recordServerError(
					"credential %q initialize carried session %q",
					credential,
					got,
				)
			}
			writer.Header().Set("Mcp-Session-Id", expectedSession)
			writeLegacyInitialize(t, writer, rpcRequest.ID)
		case "notifications/initialized":
			if got := request.Header.Get("Mcp-Session-Id"); got != expectedSession {
				recordServerError(
					"credential %q initialized session = %q, want %q",
					credential,
					got,
					expectedSession,
				)
			}
			writer.WriteHeader(http.StatusAccepted)
		case "tools/call":
			if got := request.Header.Get("Mcp-Session-Id"); got != expectedSession {
				recordServerError(
					"credential %q call session = %q, want %q",
					credential,
					got,
					expectedSession,
				)
			}
			var params struct {
				Arguments map[string]string `json:"arguments"`
			}
			if err := json.Unmarshal(rpcRequest.Params, &params); err != nil {
				recordServerError(
					"credential %q params decode: %v", credential, err)
			}
			if params.Arguments["credential"] != credential {
				recordServerError(
					"credential %q saw arguments for %q",
					credential,
					params.Arguments["credential"],
				)
			}
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"content": []map[string]any{
					{
						"type": "text",
						"text": credential + ":" + strings.Repeat("x", 512),
					},
				},
			})
		default:
			recordServerError(
				"credential %q unexpected method %q",
				credential,
				rpcRequest.Method,
			)
			writer.WriteHeader(http.StatusBadRequest)
		}
	}))
	defer server.Close()

	client, err := mcp.NewClient(mcp.Options{
		Endpoint:         server.URL + "/mcp",
		Auth:             testAuth,
		AllowedTools:     []string{"fixed_tool"},
		Transport:        server.Client().Transport,
		MaxResponseBytes: 1024,
	})
	require.NoError(t, err)

	credentials := []string{"alpha-secret", "beta-secret"}
	var (
		waitGroup sync.WaitGroup
		resultsMu sync.Mutex
		results   = make(map[string]string, len(credentials))
		callErrs  = make(chan error, len(credentials))
	)
	for _, credential := range credentials {
		credential := credential
		waitGroup.Add(1)
		go func() {
			defer waitGroup.Done()
			result, callErr := client.CallTool(
				context.Background(),
				runner.Credentials{"apiKey": credential},
				"fixed_tool",
				map[string]any{"credential": credential},
			)
			if callErr != nil {
				callErrs <- callErr
				return
			}
			if len(result.Content) != 1 {
				callErrs <- fmt.Errorf(
					"credential %q content count = %d",
					credential,
					len(result.Content),
				)
				return
			}
			var content struct {
				Text string `json:"text"`
			}
			if err := json.Unmarshal(result.Content[0], &content); err != nil {
				callErrs <- err
				return
			}
			resultsMu.Lock()
			results[credential] = content.Text
			resultsMu.Unlock()
		}()
	}
	waitGroup.Wait()
	close(callErrs)
	for callErr := range callErrs {
		require.NoError(t, callErr)
	}

	serverErrorsMu.Lock()
	defer serverErrorsMu.Unlock()
	require.Empty(t, serverErrors)
	for _, credential := range credentials {
		require.Equal(
			t,
			credential+":"+strings.Repeat("x", 512),
			results[credential],
		)
	}
}