package pkulaw

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

	"github.com/stretchr/testify/require"
)

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

type pkulawMCPWireDouble struct {
	t          *testing.T
	credential string
	path       string
	result     func(pkulawRPCRequest) any

	mu      sync.Mutex
	methods []string
}

func (server *pkulawMCPWireDouble) 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.Equal(server.t, server.path, 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.Equal(
			server.t,
			pkulawSessionID(server.path),
			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 pkulawRPCRequest
	require.NoError(server.t, json.Unmarshal(body, &rpc))
	server.record(rpc.Method)

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

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

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

func writePKULawRPCResult(
	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 writePKULawRPCError(
	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,
		},
	}))
}