// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package mcp

import (
	"context"
	"fmt"
	"log/slog"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/google/jsonschema-go/jsonschema"
)

type Item struct {
	Name  string
	Value string
}

type ListTestParams struct {
	Cursor string
}

func (p *ListTestParams) cursorPtr() *string {
	return &p.Cursor
}

type ListTestResult struct {
	Items      []*Item
	NextCursor string
}

func (r *ListTestResult) nextCursorPtr() *string {
	return &r.NextCursor
}

var allItems = []*Item{
	{"alpha", "val-A"},
	{"bravo", "val-B"},
	{"charlie", "val-C"},
	{"delta", "val-D"},
	{"echo", "val-E"},
	{"foxtrot", "val-F"},
	{"golf", "val-G"},
	{"hotel", "val-H"},
	{"india", "val-I"},
	{"juliet", "val-J"},
	{"kilo", "val-K"},
}

// generatePaginatedResults is a helper to create a sequence of mock responses for pagination.
// It simulates a server returning items in pages based on a given page size.
func generatePaginatedResults(all []*Item, pageSize int) []*ListTestResult {
	if len(all) == 0 {
		return []*ListTestResult{{Items: []*Item{}, NextCursor: ""}}
	}
	if pageSize <= 0 {
		panic("pageSize must be greater than 0")
	}
	numPages := (len(all) + pageSize - 1) / pageSize // Ceiling division
	var results []*ListTestResult
	for i := range numPages {
		startIndex := i * pageSize
		endIndex := min(startIndex+pageSize, len(all)) // Use min to prevent out of bounds
		nextCursor := ""
		if endIndex < len(all) { // If there are more items after this page
			nextCursor = fmt.Sprintf("cursor_%d", endIndex)
		}
		results = append(results, &ListTestResult{Items: all[startIndex:endIndex], NextCursor: nextCursor})
	}
	return results
}

func TestClientPaginateBasic(t *testing.T) {
	ctx := context.Background()
	testCases := []struct {
		name          string
		results       []*ListTestResult
		mockError     error
		initialParams *ListTestParams
		expected      []*Item
		expectError   bool
	}{
		{
			name:     "SinglePageAllItems",
			results:  generatePaginatedResults(allItems, len(allItems)),
			expected: allItems,
		},
		{
			name:     "MultiplePages",
			results:  generatePaginatedResults(allItems, 3),
			expected: allItems,
		},
		{
			name:     "EmptyResults",
			results:  generatePaginatedResults([]*Item{}, 10),
			expected: nil,
		},
		{
			name:        "ListFuncReturnsErrorImmediately",
			results:     []*ListTestResult{{}},
			mockError:   fmt.Errorf("API error on first call"),
			expected:    nil,
			expectError: true,
		},
		{
			name:          "InitialCursorProvided",
			initialParams: &ListTestParams{Cursor: "cursor_2"},
			results:       generatePaginatedResults(allItems[2:], 3),
			expected:      allItems[2:],
		},
		{
			name:          "CursorBeyondAllItems",
			initialParams: &ListTestParams{Cursor: "cursor_999"},
			results:       []*ListTestResult{{Items: []*Item{}, NextCursor: ""}},
			expected:      nil,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
				if len(tc.results) == 0 {
					t.Fatalf("listFunc called but no more results defined for test case %q", tc.name)
				}
				res := tc.results[0]
				tc.results = tc.results[1:]
				var err error
				if tc.mockError != nil {
					err = tc.mockError
				}
				return res, err
			}

			params := tc.initialParams
			if tc.initialParams == nil {
				params = &ListTestParams{}
			}

			var gotItems []*Item
			var iterationErr error
			seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items })
			for item, err := range seq {
				if err != nil {
					iterationErr = err
					break
				}
				gotItems = append(gotItems, item)
			}
			if tc.expectError {
				if iterationErr == nil {
					t.Errorf("paginate() expected an error during iteration, but got none")
				}
			} else {
				if iterationErr != nil {
					t.Errorf("paginate() got: %v, want: nil", iterationErr)
				}
			}
			if diff := cmp.Diff(tc.expected, gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
				t.Fatalf("paginate() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func TestClientLogger(t *testing.T) {
	// Case 1: No logger provided
	c1 := NewClient(&Implementation{Name: "test", Version: "1.0"}, nil)
	if c1.opts.Logger == nil {
		t.Error("expected default logger, got nil")
	}

	// Case 2: Logger provided
	logger := slog.Default()
	c2 := NewClient(&Implementation{Name: "test", Version: "1.0"}, &ClientOptions{
		Logger: logger,
	})
	if c2.opts.Logger != logger {
		t.Error("expected provided logger, got different one")
	}
}

func TestClientPaginateVariousPageSizes(t *testing.T) {
	ctx := context.Background()
	for i := 1; i < len(allItems)+1; i++ {
		testname := fmt.Sprintf("PageSize=%d", i)
		t.Run(testname, func(t *testing.T) {
			results := generatePaginatedResults(allItems, i)
			listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
				res := results[0]
				results = results[1:]
				return res, nil
			}
			var gotItems []*Item
			seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items })
			for item, err := range seq {
				if err != nil {
					t.Fatalf("paginate() unexpected error during iteration: %v", err)
				}
				gotItems = append(gotItems, item)
			}
			if diff := cmp.Diff(allItems, gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
				t.Fatalf("paginate() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func TestClientCapabilities(t *testing.T) {
	testCases := []struct {
		name             string
		configureClient  func(s *Client)
		clientOpts       ClientOptions
		protocolVersion  string // defaults to latestProtocolVersion if empty
		wantCapabilities *ClientCapabilities
	}{
		{
			name:            "default",
			configureClient: func(s *Client) {},
			wantCapabilities: &ClientCapabilities{
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
			},
		},
		{
			name:            "with sampling",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
					return nil, nil
				},
			},
			wantCapabilities: &ClientCapabilities{
				Roots:    RootCapabilities{ListChanged: true},
				RootsV2:  &RootCapabilities{ListChanged: true},
				Sampling: &SamplingCapabilities{},
			},
		},
		{
			name:            "with elicitation",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
					return nil, nil
				},
			},
			protocolVersion: protocolVersion20251125,
			wantCapabilities: &ClientCapabilities{
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
				Elicitation: &ElicitationCapabilities{
					Form: &FormElicitationCapabilities{},
				},
			},
		},
		{
			name:            "with elicitation (old protocol)",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
					return nil, nil
				},
			},
			protocolVersion: protocolVersion20250618,
			wantCapabilities: &ClientCapabilities{
				Roots:       RootCapabilities{ListChanged: true},
				RootsV2:     &RootCapabilities{ListChanged: true},
				Elicitation: &ElicitationCapabilities{},
			},
		},
		{
			name:            "with URL elicitation",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					Roots:   RootCapabilities{ListChanged: true},
					RootsV2: &RootCapabilities{ListChanged: true},
					Elicitation: &ElicitationCapabilities{
						URL: &URLElicitationCapabilities{},
					},
				},
				ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
					return nil, nil
				},
			},
			wantCapabilities: &ClientCapabilities{
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
				Elicitation: &ElicitationCapabilities{
					URL: &URLElicitationCapabilities{},
				},
			},
		},
		{
			name:            "with form and URL elicitation",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					Roots:   RootCapabilities{ListChanged: true},
					RootsV2: &RootCapabilities{ListChanged: true},
					Elicitation: &ElicitationCapabilities{
						Form: &FormElicitationCapabilities{},
						URL:  &URLElicitationCapabilities{},
					},
				},
				ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
					return nil, nil
				},
			},
			wantCapabilities: &ClientCapabilities{
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
				Elicitation: &ElicitationCapabilities{
					Form: &FormElicitationCapabilities{},
					URL:  &URLElicitationCapabilities{},
				},
			},
		},
		{
			name:            "no capabilities",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{},
			},
			wantCapabilities: &ClientCapabilities{},
		},
		{
			name:            "no roots",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					Sampling: &SamplingCapabilities{},
				},
			},
			wantCapabilities: &ClientCapabilities{
				Sampling: &SamplingCapabilities{},
			},
		},
		{
			name:            "roots-no list",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					RootsV2: &RootCapabilities{ListChanged: false},
				},
			},
			wantCapabilities: &ClientCapabilities{
				RootsV2: &RootCapabilities{ListChanged: false},
			},
		},
		{
			name:            "custom capabilities with sampling",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					RootsV2: &RootCapabilities{ListChanged: true},
				},
				CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
					return nil, nil
				},
			},
			wantCapabilities: &ClientCapabilities{
				Roots:    RootCapabilities{ListChanged: true},
				RootsV2:  &RootCapabilities{ListChanged: true},
				Sampling: &SamplingCapabilities{},
			},
		},
		{
			name:            "elicitation override",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					Elicitation: &ElicitationCapabilities{
						URL: &URLElicitationCapabilities{},
					},
				},
				ElicitationHandler: func(context.Context, *ElicitRequest) (*ElicitResult, error) {
					return nil, nil
				},
			},
			wantCapabilities: &ClientCapabilities{
				Elicitation: &ElicitationCapabilities{
					URL: &URLElicitationCapabilities{},
				},
			},
		},
		{
			name:            "custom capabilities with experimental",
			configureClient: func(s *Client) {},
			clientOpts: ClientOptions{
				Capabilities: &ClientCapabilities{
					Experimental: map[string]any{"custom": "value"},
					RootsV2:      &RootCapabilities{ListChanged: true},
				},
			},
			wantCapabilities: &ClientCapabilities{
				Experimental: map[string]any{"custom": "value"},
				Roots:        RootCapabilities{ListChanged: true},
				RootsV2:      &RootCapabilities{ListChanged: true},
			},
		},
		{
			name:            "extensions preserved",
			configureClient: func(s *Client) {},
			clientOpts: func() ClientOptions {
				caps := &ClientCapabilities{
					RootsV2: &RootCapabilities{ListChanged: true},
				}
				caps.AddExtension("io.example/ext1", map[string]any{"key": "value"})
				caps.AddExtension("io.example/ext2", nil)
				return ClientOptions{Capabilities: caps}
			}(),
			wantCapabilities: &ClientCapabilities{
				Extensions: map[string]any{
					"io.example/ext1": map[string]any{"key": "value"},
					"io.example/ext2": map[string]any{},
				},
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			client := NewClient(testImpl, &tc.clientOpts)
			tc.configureClient(client)
			protocolVersion := tc.protocolVersion
			if protocolVersion == "" {
				protocolVersion = latestProtocolVersion
			}
			gotCapabilities := client.capabilities(protocolVersion)
			if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" {
				t.Errorf("capabilities() mismatch (-want +got):\n%s", diff)
			}
		})
	}
}

func TestClientCapabilitiesOverWire(t *testing.T) {
	testCases := []struct {
		name             string
		clientOpts       *ClientOptions
		wantCapabilities *ClientCapabilities
	}{
		{
			name:       "Default capabilities",
			clientOpts: nil,
			wantCapabilities: &ClientCapabilities{
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
			},
		},
		{
			name: "Custom Capabilities with roots listChanged false",
			clientOpts: &ClientOptions{
				Capabilities: &ClientCapabilities{
					RootsV2: &RootCapabilities{ListChanged: false},
				},
			},
			wantCapabilities: &ClientCapabilities{
				Roots:   RootCapabilities{ListChanged: false},
				RootsV2: &RootCapabilities{ListChanged: false},
			},
		},
		{
			name: "Dynamic sampling capability",
			clientOpts: &ClientOptions{
				Capabilities: &ClientCapabilities{
					RootsV2: &RootCapabilities{ListChanged: true},
				},
				CreateMessageHandler: func(context.Context, *CreateMessageRequest) (*CreateMessageResult, error) {
					return nil, nil
				},
			},
			wantCapabilities: &ClientCapabilities{
				Roots:    RootCapabilities{ListChanged: true},
				RootsV2:  &RootCapabilities{ListChanged: true},
				Sampling: &SamplingCapabilities{},
			},
		},
		{
			name: "Empty capabilities disables defaults",
			clientOpts: &ClientOptions{
				Capabilities: &ClientCapabilities{},
			},
			wantCapabilities: &ClientCapabilities{},
		},
		{
			name: "Extensions over wire",
			clientOpts: func() *ClientOptions {
				caps := &ClientCapabilities{
					RootsV2: &RootCapabilities{ListChanged: true},
				}
				caps.AddExtension("io.example/ext", map[string]any{"key": "value"})
				return &ClientOptions{Capabilities: caps}
			}(),
			wantCapabilities: &ClientCapabilities{
				Extensions: map[string]any{
					"io.example/ext": map[string]any{"key": "value"},
				},
				Roots:   RootCapabilities{ListChanged: true},
				RootsV2: &RootCapabilities{ListChanged: true},
			},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			ctx := context.Background()

			// Create client.
			impl := &Implementation{Name: "testClient", Version: "v1.0.0"}
			client := NewClient(impl, tc.clientOpts)

			// Connect client and server.
			cTransport, sTransport := NewInMemoryTransports()
			server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
			ss, err := server.Connect(ctx, sTransport, nil)
			if err != nil {
				t.Fatal(err)
			}
			defer ss.Close()

			cs, err := client.Connect(ctx, cTransport, nil)
			if err != nil {
				t.Fatal(err)
			}
			defer cs.Close()

			// Check that the server received the expected capabilities.
			initParams := ss.InitializeParams()
			if initParams == nil {
				t.Fatal("InitializeParams is nil")
			}

			if diff := cmp.Diff(tc.wantCapabilities, initParams.Capabilities); diff != "" {
				t.Errorf("Capabilities mismatch (-want +got):\n%s", diff)
			}
		})
	}
}