package pkulaw

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

	"github.com/stretchr/testify/require"

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

type pkulawProbeWireDouble struct {
	t          *testing.T
	credential string
	failures   map[string]int

	mu      sync.Mutex
	methods map[string][]string
}

func (server *pkulawProbeWireDouble) RoundTrip(
	request *http.Request,
) (*http.Response, error) {
	server.t.Helper()
	require.Equal(server.t, "https", request.URL.Scheme)
	require.Equal(server.t, "apim-gateway.pkulaw.com", request.URL.Host)
	require.Contains(server.t, endpointPaths(), request.URL.Path)
	require.Equal(
		server.t,
		"Bearer "+server.credential,
		request.Header.Get("Authorization"),
	)

	path := request.URL.Path
	recorder := httptest.NewRecorder()
	if request.Method == http.MethodDelete {
		require.Equal(
			server.t,
			pkulawSessionID(path),
			request.Header.Get("Mcp-Session-Id"),
		)
		server.record(path, "DELETE")
		recorder.WriteHeader(http.StatusNoContent)
		return recorder.Result(), nil
	}
	require.Equal(server.t, http.MethodPost, request.Method)
	body, err := io.ReadAll(request.Body)
	require.NoError(server.t, err)
	var rpc pkulawRPCRequest
	require.NoError(server.t, json.Unmarshal(body, &rpc))
	server.record(path, rpc.Method)

	switch rpc.Method {
	case "server/discover":
		writePKULawRPCError(
			server.t,
			recorder,
			rpc.ID,
			-32601,
			"legacy only",
		)
	case "initialize":
		if status := server.failures[path]; status != 0 {
			recorder.WriteHeader(status)
			_, err := recorder.WriteString(
				"upstream body contains " + server.credential,
			)
			require.NoError(server.t, err)
			return recorder.Result(), nil
		}
		recorder.Header().Set("Mcp-Session-Id", pkulawSessionID(path))
		writePKULawRPCResult(server.t, recorder, rpc.ID, map[string]any{
			"protocolVersion": "2025-06-18",
			"capabilities":    map[string]any{},
			"serverInfo": map[string]any{
				"name":    "pkulaw-probe-double",
				"version": "1.0.0",
			},
		})
	case "notifications/initialized":
		require.Equal(
			server.t,
			pkulawSessionID(path),
			request.Header.Get("Mcp-Session-Id"),
		)
		recorder.WriteHeader(http.StatusAccepted)
	default:
		require.Failf(
			server.t,
			"probe issued a capability operation",
			"path=%q method=%q",
			path,
			rpc.Method,
		)
		recorder.WriteHeader(http.StatusBadRequest)
	}
	return recorder.Result(), nil
}

func (server *pkulawProbeWireDouble) record(path, method string) {
	server.mu.Lock()
	defer server.mu.Unlock()
	if server.methods == nil {
		server.methods = make(map[string][]string)
	}
	server.methods[path] = append(server.methods[path], method)
}

func (server *pkulawProbeWireDouble) recorded(path string) []string {
	server.mu.Lock()
	defer server.mu.Unlock()
	return append([]string(nil), server.methods[path]...)
}

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

	server := &pkulawProbeWireDouble{
		t:          t,
		credential: "pkulaw-probe-secret",
		failures:   map[string]int{},
	}
	clients, err := newProbeClients(server)
	require.NoError(t, err)
	config, err := prepareConnection(
		context.Background(),
		clients,
		runner.Credentials{"apiKey": "pkulaw-probe-secret"},
	)
	require.NoError(t, err)
	require.Equal(t, runner.ConnectionConfig{
		"topology": TopologyAggregate,
	}, config)
	require.Equal(t, []string{
		"server/discover",
		"initialize",
		"notifications/initialized",
		"DELETE",
	}, server.recorded("/mcp-law-agg/1.0.0/mcp"))
	for _, path := range fallbackEndpointPaths() {
		require.Empty(t, server.recorded(path))
	}
}

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

	server := &pkulawProbeWireDouble{
		t:          t,
		credential: "pkulaw-probe-secret",
		failures: map[string]int{
			"/mcp-law-agg/1.0.0/mcp": http.StatusNotFound,
		},
	}
	clients, err := newProbeClients(server)
	require.NoError(t, err)
	config, err := prepareConnection(
		context.Background(),
		clients,
		runner.Credentials{"apiKey": "pkulaw-probe-secret"},
	)
	require.NoError(t, err)
	require.Equal(t, runner.ConnectionConfig{
		"topology": TopologyFallback,
	}, config)
	for _, path := range fallbackEndpointPaths() {
		require.Equal(t, []string{
			"server/discover",
			"initialize",
			"notifications/initialized",
			"DELETE",
		}, server.recorded(path))
	}
}

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

	server := &pkulawProbeWireDouble{
		t:          t,
		credential: "pkulaw-probe-secret",
		failures: map[string]int{
			"/mcp-law-agg/1.0.0/mcp":       http.StatusUnauthorized,
			"/mcp-case-search-service/mcp": http.StatusForbidden,
		},
	}
	clients, err := newProbeClients(server)
	require.NoError(t, err)
	config, err := prepareConnection(
		context.Background(),
		clients,
		runner.Credentials{"apiKey": "pkulaw-probe-secret"},
	)
	require.Nil(t, config)
	require.Error(t, err)
	require.Equal(
		t,
		runner.ErrorProviderAuthUnavailable,
		runner.CodeOf(err),
	)
	for _, path := range fallbackEndpointPaths() {
		require.Contains(t, server.recorded(path), "initialize")
	}
}

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

	const credential = "pkulaw-aggregate-auth-secret"
	failures := map[string]int{
		"/mcp-law-agg/1.0.0/mcp": http.StatusUnauthorized,
	}
	for _, path := range fallbackEndpointPaths() {
		failures[path] = http.StatusNotFound
	}
	server := &pkulawProbeWireDouble{
		t:          t,
		credential: credential,
		failures:   failures,
	}
	clients, err := newProbeClients(server)
	require.NoError(t, err)

	config, err := prepareConnection(
		context.Background(),
		clients,
		runner.Credentials{"apiKey": credential},
	)
	require.Nil(t, config)
	require.Error(t, err)
	require.Equal(
		t,
		runner.ErrorProviderAuthUnavailable,
		runner.CodeOf(err),
	)
	require.NotContains(t, err.Error(), credential)
	require.NotContains(t, err.Error(), "upstream body")
}

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

	var calls atomic.Int64
	clients, err := newProbeClients(roundTripperFunc(func(
		*http.Request,
	) (*http.Response, error) {
		calls.Add(1)
		return nil, context.Canceled
	}))
	require.NoError(t, err)
	client := clients[EndpointAggregate]
	credentials := runner.Credentials{"apiKey": "pkulaw-probe-secret"}

	_, err = client.CallTool(
		context.Background(),
		credentials,
		"search_article",
		map[string]any{"query": "民法典"},
	)
	require.Error(t, err)
	require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))

	_, err = client.ReadResource(
		context.Background(),
		credentials,
		"pkulaw://anything",
	)
	require.Error(t, err)
	require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))
	require.Zero(t, calls.Load())
}

func endpointPaths() []string {
	paths := make([]string, 0, len(endpointPolicies))
	for _, policy := range endpointPolicies {
		paths = append(paths, strings.TrimPrefix(policy.Endpoint, Gateway))
	}
	return paths
}

func fallbackEndpointPaths() []string {
	return []string{
		"/mcp-law-search-service/mcp",
		"/mcp-law/mcp",
		"/mcp-case-search-service/mcp",
		"/mcp-case/mcp",
	}
}

func pkulawSessionID(path string) string {
	replacer := map[rune]rune{'/': '-', '_': '-'}
	result := make([]rune, 0, len(path)+8)
	result = append(result, []rune("pkulaw")...)
	for _, character := range path {
		if replacement, ok := replacer[character]; ok {
			result = append(result, replacement)
			continue
		}
		result = append(result, character)
	}
	return string(result)
}

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

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