package mcp_test

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestOperationAllowlistsFailBeforeNetwork(t *testing.T) {
	t.Parallel()

	var calls atomic.Int64
	client, err := mcp.NewClient(mcp.Options{
		Endpoint:            "https://mcp.example.com/mcp",
		Auth:                testAuth,
		AllowedTools:        []string{"fixed_tool"},
		AllowedResourceURIs: []string{"fixed://resource"},
		Transport: roundTripperFunc(func(
			*http.Request,
		) (*http.Response, error) {
			calls.Add(1)
			return nil, errors.New("network must not be reached")
		}),
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "secret"},
		"user_selected_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))

	_, err = client.ReadResource(
		context.Background(),
		runner.Credentials{"apiKey": "secret"},
		"user://selected",
	)
	require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "secret"},
		"fixed_tool",
		[]string{"not", "an", "object"},
	)
	require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))

	err = client.Probe(
		context.Background(),
		runner.Credentials{"apiKey": "secret"},
	)
	require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))
	require.Zero(t, calls.Load())
}

func TestInvalidCredentialsFailBeforeNetwork(t *testing.T) {
	t.Parallel()

	var calls atomic.Int64
	client, err := mcp.NewClient(mcp.Options{
		Endpoint:     "https://mcp.example.com/mcp",
		Auth:         testAuth,
		AllowedTools: []string{"fixed_tool"},
		Transport: roundTripperFunc(func(
			*http.Request,
		) (*http.Response, error) {
			calls.Add(1)
			return nil, errors.New("network must not be reached")
		}),
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorProviderAuthUnavailable,
		runner.CodeOf(err))
	requireRuntimeErrorHasNoRawCause(t, err, "apiKey")
	require.Zero(t, calls.Load())
}

func TestConnectFailureWithSessionPerformsDeleteCleanup(t *testing.T) {
	t.Parallel()

	const sessionID = "unsupported-version-session"
	deleteSeen := make(chan struct{}, 1)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			if request.Header.Get("Mcp-Session-Id") != sessionID {
				t.Errorf("DELETE session ID = %q",
					request.Header.Get("Mcp-Session-Id"))
			}
			if got := request.Header.Get("Mcp-Protocol-Version"); got != "" {
				t.Errorf("DELETE retained stale protocol version %q", got)
			}
			if request.Header.Get("X-MCP-Key") != "cleanup-secret" {
				t.Errorf("DELETE missed invocation credential")
			}
			deleteSeen <- struct{}{}
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writeRPCError(t, writer, rpcRequest.ID, -32601, "legacy")
		case "initialize":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"protocolVersion": "2099-01-01",
				"capabilities":    map[string]any{},
				"serverInfo": map[string]any{
					"name":    "unsupported-double",
					"version": "1.0.0",
				},
			})
		default:
			t.Errorf("unexpected method %q", 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,
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "cleanup-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	requireRuntimeErrorHasNoRawCause(t, err, "2099-01-01")
	select {
	case <-deleteSeen:
	case <-time.After(time.Second):
		t.Fatal("Connect failure did not issue DELETE cleanup")
	}
}

func TestCleanupBypassesExhaustedResponseBudget(t *testing.T) {
	t.Parallel()

	const sessionID = "budget-session"
	deleteSeen := make(chan struct{}, 1)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			if request.Header.Get("Mcp-Session-Id") != sessionID {
				t.Errorf("DELETE session ID = %q",
					request.Header.Get("Mcp-Session-Id"))
			}
			deleteSeen <- struct{}{}
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		writer.Header().Set("Mcp-Session-Id", sessionID)
		writer.Header().Set("Content-Type", "application/json")
		_, err := writer.Write(append(
			mustJSON(t, map[string]any{
				"jsonrpc": "2.0",
				"id":      json.RawMessage(rpcRequest.ID),
				"result": map[string]any{
					"supportedVersions": []string{"2026-07-28"},
					"capabilities":      map[string]any{},
					"instructions":      string(make([]byte, 512)),
				},
			}),
			'\n',
		))
		if err != nil {
			t.Errorf("write oversized response: %v", err)
		}
	}))
	defer server.Close()

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

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "budget-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	select {
	case <-deleteSeen:
	case <-time.After(time.Second):
		t.Fatal("exhausted invocation budget suppressed DELETE cleanup")
	}
}

func TestRequestBudgetIsCumulativeAndReservesCleanup(t *testing.T) {
	t.Parallel()

	const sessionID = "request-budget-session"
	var toolReached atomic.Bool
	deleteSeen := make(chan struct{}, 1)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			deleteSeen <- struct{}{}
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeModernDiscover(t, writer, rpcRequest.ID)
		case "tools/call":
			toolReached.Store(true)
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"content": []map[string]any{
					{"type": "text", "text": "unexpected"},
				},
			})
		default:
			t.Errorf("unexpected method %q", 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,
		MaxRequestBytes: 512,
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "request-budget-secret"},
		"fixed_tool",
		map[string]any{"payload": strings.Repeat("x", 300)},
	)
	require.Equal(t, runner.ErrorValidation,
		runner.CodeOf(err))
	require.False(t, toolReached.Load(),
		"over-budget call reached the network")
	select {
	case <-deleteSeen:
	case <-time.After(time.Second):
		t.Fatal("request budget exhaustion suppressed DELETE cleanup")
	}
}

func TestResponseBudgetIsCumulativeAcrossHandshakeAndCall(t *testing.T) {
	t.Parallel()

	const sessionID = "response-budget-session"
	deleteSeen := make(chan struct{}, 1)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			deleteSeen <- struct{}{}
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeModernDiscover(t, writer, rpcRequest.ID)
		case "tools/call":
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"content": []map[string]any{
					{
						"type": "text",
						"text": string(make([]byte, 512)),
					},
				},
			})
		default:
			t.Errorf("unexpected method %q", 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: 512,
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "response-budget-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	select {
	case <-deleteSeen:
	case <-time.After(time.Second):
		t.Fatal("response budget exhaustion suppressed DELETE cleanup")
	}
}

func TestTimeoutBoundsWholeHandshakeAndCall(t *testing.T) {
	const (
		sessionID      = "total-timeout-session"
		totalTimeout   = 150 * time.Millisecond
		maxTestRuntime = 500 * time.Millisecond
	)
	deleteSeen := make(chan struct{}, 1)
	toolSeen := make(chan struct{}, 1)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			deleteSeen <- struct{}{}
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeModernDiscover(t, writer, rpcRequest.ID)
		case "tools/call":
			toolSeen <- struct{}{}
			<-request.Context().Done()
		case "notifications/cancelled":
			writer.WriteHeader(http.StatusAccepted)
		default:
			t.Errorf("unexpected method %q", 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,
		Timeout:        totalTimeout,
		CleanupTimeout: time.Second,
	})
	require.NoError(t, err)

	start := time.Now()
	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "total-timeout-secret"},
		"fixed_tool",
		map[string]any{},
	)
	elapsed := time.Since(start)
	require.ErrorIs(t, err, context.DeadlineExceeded)
	require.Less(t, elapsed, maxTestRuntime)
	select {
	case <-toolSeen:
	default:
		t.Fatal("timeout fired before the operation started")
	}
	select {
	case <-deleteSeen:
	case <-time.After(time.Second):
		t.Fatal("total timeout prevented DELETE cleanup")
	}
}

func TestTimeoutBoundsDetachedCancellationNotification(t *testing.T) {
	const (
		sessionID      = "detached-cancel-session"
		totalTimeout   = 50 * time.Millisecond
		releaseDelay   = 500 * time.Millisecond
		maxTestRuntime = 250 * time.Millisecond
	)
	releaseCancellation := make(chan struct{})
	time.AfterFunc(releaseDelay, func() {
		close(releaseCancellation)
	})
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeModernDiscover(t, writer, rpcRequest.ID)
		case "tools/call":
			<-request.Context().Done()
		case "notifications/cancelled":
			select {
			case <-releaseCancellation:
				writer.WriteHeader(http.StatusAccepted)
			case <-request.Context().Done():
			}
		default:
			t.Errorf("unexpected method %q", 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,
		Timeout:        totalTimeout,
		CleanupTimeout: time.Second,
	})
	require.NoError(t, err)

	start := time.Now()
	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "deadline-secret"},
		"fixed_tool",
		map[string]any{},
	)
	elapsed := time.Since(start)
	require.ErrorIs(t, err, context.DeadlineExceeded)
	require.Less(t, elapsed, maxTestRuntime,
		"detached SDK cancellation POST exceeded the invocation timeout")
}

func TestCleanupStatusDoesNotDiscardSuccessfulResult(t *testing.T) {
	t.Parallel()

	for _, cleanupStatus := range []int{
		http.StatusNotFound,
		http.StatusGone,
		http.StatusInternalServerError,
	} {
		cleanupStatus := cleanupStatus
		t.Run(http.StatusText(cleanupStatus), func(t *testing.T) {
			t.Parallel()
			const sessionID = "cleanup-status-session"
			server := httptest.NewTLSServer(http.HandlerFunc(func(
				writer http.ResponseWriter,
				request *http.Request,
			) {
				if request.Method == http.MethodDelete {
					writer.WriteHeader(cleanupStatus)
					_, _ = writer.Write([]byte("raw-cleanup-secret"))
					return
				}
				rpcRequest, ok := decodeRPCRequest(t, request)
				if !ok {
					writer.WriteHeader(http.StatusBadRequest)
					return
				}
				switch rpcRequest.Method {
				case "server/discover":
					writeRPCError(
						t, writer, rpcRequest.ID, -32601, "legacy")
				case "initialize":
					writer.Header().Set("Mcp-Session-Id", sessionID)
					writeLegacyInitialize(t, writer, rpcRequest.ID)
				case "notifications/initialized":
					writer.WriteHeader(http.StatusAccepted)
				case "tools/call":
					writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
						"content": []map[string]any{
							{"type": "text", "text": "ok"},
						},
					})
				default:
					t.Errorf("unexpected method %q", 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,
			})
			require.NoError(t, err)

			result, err := client.CallTool(
				context.Background(),
				runner.Credentials{
					"apiKey": "cleanup-status-secret",
				},
				"fixed_tool",
				map[string]any{},
			)
			require.NoError(t, err)
			require.Len(t, result.Content, 1)
			require.JSONEq(
				t,
				`{"type":"text","text":"ok"}`,
				string(result.Content[0]),
			)
		})
	}
}

func TestRedirectIsNotFollowed(t *testing.T) {
	t.Parallel()

	var redirected atomic.Bool
	alternate := httptest.NewTLSServer(http.HandlerFunc(func(
		http.ResponseWriter,
		*http.Request,
	) {
		redirected.Store(true)
	}))
	defer alternate.Close()

	const sessionID = "redirect-session"
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method == http.MethodDelete {
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writeRPCError(t, writer, rpcRequest.ID, -32601, "legacy")
		case "initialize":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeLegacyInitialize(t, writer, rpcRequest.ID)
		case "notifications/initialized":
			writer.WriteHeader(http.StatusAccepted)
		case "tools/call":
			http.Redirect(
				writer,
				request,
				alternate.URL+"/stolen",
				http.StatusTemporaryRedirect,
			)
		default:
			t.Errorf("unexpected method %q", 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,
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "redirect-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	require.False(t, redirected.Load())
	requireRuntimeErrorHasNoRawCause(t, err, alternate.URL)
}

type rpcRequest struct {
	ID     json.RawMessage `json:"id"`
	Method string          `json:"method"`
	Params json.RawMessage `json:"params"`
}

func decodeRPCRequest(
	t *testing.T,
	request *http.Request,
) (rpcRequest, bool) {
	t.Helper()
	body, err := io.ReadAll(request.Body)
	if err != nil {
		t.Errorf("read request: %v", err)
		return rpcRequest{}, false
	}
	var decoded rpcRequest
	if err := json.Unmarshal(body, &decoded); err != nil {
		t.Errorf("decode request: %v", err)
		return rpcRequest{}, false
	}
	return decoded, true
}

func writeLegacyInitialize(
	t *testing.T,
	writer http.ResponseWriter,
	id json.RawMessage,
) {
	t.Helper()
	writeRPCResult(t, writer, id, map[string]any{
		"protocolVersion": "2025-11-25",
		"capabilities":    map[string]any{},
		"serverInfo": map[string]any{
			"name":    "legacy-double",
			"version": "1.0.0",
		},
	})
}

func requireRuntimeErrorHasNoRawCause(
	t *testing.T,
	err error,
	raw string,
) {
	t.Helper()
	var runtimeError *runner.Error
	require.ErrorAs(t, err, &runtimeError)
	require.Nil(t, runtimeError.Cause)
	require.NotContains(t, runtimeError.Message, raw)
}

type roundTripperFunc func(*http.Request) (*http.Response, error)

func (function roundTripperFunc) RoundTrip(
	request *http.Request,
) (*http.Response, error) {
	return function(request)
}