package jin10

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

	"github.com/stretchr/testify/require"

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

type staticCredentialSource struct {
	credentials runner.Credentials
	err         error
}

func (source staticCredentialSource) ResolveProviderCredentials(
	context.Context,
	string,
	string,
) (runner.Credentials, error) {
	if source.err != nil {
		return nil, source.err
	}
	return source.credentials.Clone(), nil
}

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

type mcpWireDouble struct {
	t               *testing.T
	credential      string
	operationResult func(rpcRequest) any

	mu      sync.Mutex
	methods []string
}

func (server *mcpWireDouble) RoundTrip(
	request *http.Request,
) (*http.Response, error) {
	server.t.Helper()
	require.Equal(server.t, "https", request.URL.Scheme)
	require.Equal(server.t, "mcp.jin10.com", request.URL.Host)
	require.Equal(server.t, "/mcp", request.URL.Path)
	require.Empty(server.t, request.URL.RawQuery)
	require.Equal(
		server.t,
		"Bearer "+server.credential,
		request.Header.Get("Authorization"),
	)

	recorder := httptest.NewRecorder()
	if request.Method == http.MethodDelete {
		require.NotEmpty(server.t, request.Header.Get("Mcp-Session-Id"))
		server.record("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 rpcRequest
	require.NoError(server.t, json.Unmarshal(body, &rpc))
	server.record(rpc.Method)

	switch rpc.Method {
	case "server/discover":
		writeRPCError(server.t, recorder, rpc.ID, -32601, "legacy only")
	case "initialize":
		recorder.Header().Set("Mcp-Session-Id", "jin10-test-session")
		writeRPCResult(server.t, recorder, rpc.ID, map[string]any{
			"protocolVersion": "2025-06-18",
			"capabilities":    map[string]any{},
			"serverInfo": map[string]any{
				"name":    "jin10-double",
				"version": "1.0.0",
			},
		})
	case "notifications/initialized":
		require.Equal(
			server.t,
			"jin10-test-session",
			request.Header.Get("Mcp-Session-Id"),
		)
		recorder.WriteHeader(http.StatusAccepted)
	case "tools/call", "resources/read":
		require.Equal(
			server.t,
			"jin10-test-session",
			request.Header.Get("Mcp-Session-Id"),
		)
		require.NotNil(server.t, server.operationResult)
		writeRPCResult(
			server.t,
			recorder,
			rpc.ID,
			server.operationResult(rpc),
		)
	default:
		require.Failf(
			server.t,
			"unexpected MCP method",
			"method=%q",
			rpc.Method,
		)
		recorder.WriteHeader(http.StatusBadRequest)
	}
	return recorder.Result(), nil
}

func (server *mcpWireDouble) record(method string) {
	server.mu.Lock()
	defer server.mu.Unlock()
	server.methods = append(server.methods, method)
}

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

func writeRPCResult(
	t *testing.T,
	recorder *httptest.ResponseRecorder,
	id json.RawMessage,
	result any,
) {
	t.Helper()
	recorder.Header().Set("Content-Type", "application/json")
	require.NoError(t, json.NewEncoder(recorder).Encode(map[string]any{
		"jsonrpc": "2.0",
		"id":      json.RawMessage(id),
		"result":  result,
	}))
}

func writeRPCError(
	t *testing.T,
	recorder *httptest.ResponseRecorder,
	id json.RawMessage,
	code int64,
	message string,
) {
	t.Helper()
	recorder.Header().Set("Content-Type", "application/json")
	require.NoError(t, json.NewEncoder(recorder).Encode(map[string]any{
		"jsonrpc": "2.0",
		"id":      json.RawMessage(id),
		"error": map[string]any{
			"code":    code,
			"message": message,
		},
	}))
}

func mustJSON(t *testing.T, value any) json.RawMessage {
	t.Helper()
	encoded, err := json.Marshal(value)
	require.NoError(t, err)
	return encoded
}