// 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 (
	"bytes"
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"slices"
	"strings"
	"testing"
	"time"
)

func TestScanEvents(t *testing.T) {
	tests := []struct {
		name    string
		input   string
		want    []Event
		wantErr string
	}{
		{
			name:  "simple event",
			input: "event: message\nid: 1\ndata: hello\n\n",
			want: []Event{
				{Name: "message", ID: "1", Data: []byte("hello")},
			},
		},
		{
			name:  "multiple data lines",
			input: "data: line 1\ndata: line 2\n\n",
			want: []Event{
				{Data: []byte("line 1\nline 2")},
			},
		},
		{
			name:  "multiple events",
			input: "data: first\n\nevent: second\ndata: second\n\n",
			want: []Event{
				{Data: []byte("first")},
				{Name: "second", Data: []byte("second")},
			},
		},
		{
			name:  "no trailing newline",
			input: "data: hello",
			want: []Event{
				{Data: []byte("hello")},
			},
		},
		{
			name:    "malformed line",
			input:   "invalid line\n\n",
			wantErr: "malformed line",
		},
		{
			name:  "message with 2 data lines and another event",
			input: "event: message\ndata: hello\ndata: hello\ndata: hello\n\nevent:keepalive",
			want: []Event{
				{Name: "message", Data: []byte("hello\nhello\nhello")},
				{Name: "keepalive"},
			},
		},
		{
			name:  "event with multiple lines",
			input: "event: message\ndata: hello\ndata: hello\ndata: hello\nid:1",
			want: []Event{
				{Name: "message", ID: "1", Data: []byte("hello\nhello\nhello")},
			},
		},
		{
			name: "multiple events, out of order keys",
			input: strings.Join([]string{
				"event:message",
				"data: hello0",
				"\n",
				"data: hello1",
				"data: hello1",
				"id:1",
				"event:message",
				"\n",
				"event:message",
				"data: hello3",
				"data: hello3",
				"id:3",
				"\n",
				"data: hello4",
				"data: hello4",
				"id:4",
				"event:message",
			}, "\n"),
			want: []Event{
				{Name: "message", Data: []byte("hello0")},
				{Name: "message", ID: "1", Data: []byte("hello1\nhello1")},
				{Name: "message", ID: "3", Data: []byte("hello3\nhello3")},
				{Name: "message", ID: "4", Data: []byte("hello4\nhello4")},
			},
		},
		{
			name:  "non-continuous data items in the event",
			input: "event: foo\ndata: 123\nretry: 5\ndata: 456",
			want: []Event{
				{Name: "foo", Data: []byte("123\n456"), Retry: "5"},
			},
		},
		{
			name:  "no-data events",
			input: "event: foo\n\nevent: bar",
			want: []Event{
				{Name: "foo"},
				{Name: "bar"},
			},
		},
		{
			name:  "empty data event",
			input: "event: foo\ndata:\n\nevent: bar",
			want: []Event{
				{Name: "foo"},
				{Name: "bar"},
			},
		},
		{

			name:    "malformed data event",
			input:   "someline",
			wantErr: "malformed event",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			r := strings.NewReader(tt.input)
			var got []Event
			var err error
			for e, err2 := range scanEventsLimited(r, DefaultMaxEventSize) {
				if err2 != nil {
					err = err2
					break
				}
				got = append(got, e)
			}

			if tt.wantErr != "" {
				if err == nil {
					t.Fatalf("scanEvents() got nil error, want error containing %q", tt.wantErr)
				}
				if !strings.Contains(err.Error(), tt.wantErr) {
					t.Fatalf("scanEvents() error = %q, want containing %q", err, tt.wantErr)
				}
				return
			}

			if err != nil {
				t.Fatalf("scanEvents() returned unexpected error: %v", err)
			}

			if len(got) != len(tt.want) {
				t.Fatalf("scanEvents() got %d events, want %d", len(got), len(tt.want))
			}

			for i := range got {
				if g, w := got[i].Name, tt.want[i].Name; g != w {
					t.Errorf("event %d: name = %q, want %q", i, g, w)
				}
				if g, w := got[i].ID, tt.want[i].ID; g != w {
					t.Errorf("event %d: id = %q, want %q", i, g, w)
				}
				if g, w := string(got[i].Data), string(tt.want[i].Data); g != w {
					t.Errorf("event %d: data = %q, want %q", i, g, w)
				}
			}
		})
	}
}

// endlessReader streams a fixed prefix once, then repeats fill forever.
type endlessReader struct {
	prefix   string
	sent     bool
	repeat   string
	repIndex int
}

func (r *endlessReader) Read(p []byte) (int, error) {
	if !r.sent {
		n := copy(p, r.prefix)
		r.sent = true
		return n, nil
	}
	for i := range p {
		p[i] = r.repeat[r.repIndex%len(r.repeat)]
		r.repIndex++
	}
	return len(p), nil
}

func (r *endlessReader) Close() error { return nil }

// TestScanEventsMaxEventSize verifies that scanEvents bounds the bytes
// buffered for a single event.
func TestScanEventsMaxEventSize(t *testing.T) {
	wantByte := byte('A')
	eventOfSize := func(size int) []byte {
		var buf bytes.Buffer
		buf.WriteString("data: ")
		buf.Write(bytes.Repeat([]byte{wantByte}, size))
		buf.WriteString("\n\n")
		return buf.Bytes()
	}

	tests := []struct {
		name            string
		reader          io.Reader
		maxEventSize    int
		wantErr         bool
		wantDataLengths []int
	}{
		{
			name:            "negative budget disables the cap",
			reader:          bytes.NewReader(eventOfSize(512)),
			maxEventSize:    -1,
			wantDataLengths: []int{512},
		},
		{
			name:            "long line under cap", // bufio.Reader default buf size is 1 << 12
			reader:          bytes.NewReader(eventOfSize(1 << 14)),
			maxEventSize:    1 << 15,
			wantDataLengths: []int{1 << 14},
		},
		{
			name:         "unbounded rejected",
			reader:       &endlessReader{prefix: "data: ", repeat: string(wantByte)},
			maxEventSize: 1024,
			wantErr:      true,
		},

		{
			name: "consecutive data fields longer than limit",
			reader: strings.NewReader(func() string {
				var b strings.Builder
				for range 2048 {
					b.WriteString("data: \n")
				}
				return b.String()
			}()),
			maxEventSize: 1024,
			wantErr:      true,
		},
		{
			name:            "budget resets between events",
			reader:          bytes.NewReader(append(eventOfSize(3*1024), eventOfSize(3*1024)...)),
			maxEventSize:    4 * 1024,
			wantDataLengths: []int{3 * 1024, 3 * 1024},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			type result struct {
				events []Event
				err    error
			}

			done := make(chan result, 1)
			go func() {
				var events []Event
				for e, err := range scanEventsLimited(tt.reader, tt.maxEventSize) {
					if err != nil {
						done <- result{events, err}
						return
					}
					events = append(events, e)
				}
				done <- result{events, nil}
			}()

			var res result
			select {
			case res = <-done:
			case <-time.After(5 * time.Second):
				t.Fatal("scanEvents did not return")
			}

			if tt.wantErr {
				if !errors.Is(res.err, errMalformedEvent) {
					t.Fatalf("got error %v, want errMalformedEvent", res.err)
				}
				return
			}
			if res.err != nil {
				t.Fatalf("unexpected error: %v", res.err)
			}
			if len(res.events) != len(tt.wantDataLengths) {
				t.Fatalf("got %d events, want %d", len(res.events), len(tt.wantDataLengths))
			}
			for i, wantLen := range tt.wantDataLengths {
				wantRepeated := string(wantByte)
				if got := res.events[i].Data; string(got) != strings.Repeat(wantRepeated, wantLen) {
					t.Errorf("event %d: got %d data bytes, want %d bytes of %s", i, len(got), wantLen, wantRepeated)
				}
			}
		})
	}
}

func TestMemoryEventStoreState(t *testing.T) {
	ctx := context.Background()

	appendEvent := func(s *MemoryEventStore, sess, stream string, data string) {
		if err := s.Append(ctx, sess, stream, []byte(data)); err != nil {
			t.Fatal(err)
		}
	}

	for _, tt := range []struct {
		name     string
		actions  func(*MemoryEventStore)
		want     string // output of debugString
		wantSize int    // value of nBytes
	}{
		{
			"appends",
			func(s *MemoryEventStore) {
				appendEvent(s, "S1", "1", "d1")
				appendEvent(s, "S1", "2", "d2")
				appendEvent(s, "S1", "1", "d3")
				appendEvent(s, "S2", "8", "d4")
			},
			"S1 1 first=0 d1 d3; S1 2 first=0 d2; S2 8 first=0 d4",
			8,
		},
		{
			"session close",
			func(s *MemoryEventStore) {
				appendEvent(s, "S1", "1", "d1")
				appendEvent(s, "S1", "2", "d2")
				appendEvent(s, "S1", "1", "d3")
				appendEvent(s, "S2", "8", "d4")
				s.SessionClosed(ctx, "S1")
			},
			"S2 8 first=0 d4",
			2,
		},
		{
			"purge",
			func(s *MemoryEventStore) {
				appendEvent(s, "S1", "1", "d1")
				appendEvent(s, "S1", "2", "d2")
				appendEvent(s, "S1", "1", "d3")
				appendEvent(s, "S2", "8", "d4")
				// We are using 8 bytes (d1,d2, d3, d4).
				// To purge 6, we remove the first of each stream, leaving only d3.
				s.SetMaxBytes(2)
			},
			// The other streams remain, because we may add to them.
			"S1 1 first=1 d3; S1 2 first=1; S2 8 first=1",
			2,
		},
		{
			"purge append",
			func(s *MemoryEventStore) {
				appendEvent(s, "S1", "1", "d1")
				appendEvent(s, "S1", "2", "d2")
				appendEvent(s, "S1", "1", "d3")
				appendEvent(s, "S2", "8", "d4")
				s.SetMaxBytes(2)
				// Up to here, identical to the "purge" case.
				// Each of these additions will result in a purge.
				appendEvent(s, "S1", "2", "d5") // remove d3
				appendEvent(s, "S1", "2", "d6") // remove d5
			},
			"S1 1 first=2; S1 2 first=2 d6; S2 8 first=1",
			2,
		},
		{
			"purge resize append",
			func(s *MemoryEventStore) {
				appendEvent(s, "S1", "1", "d1")
				appendEvent(s, "S1", "2", "d2")
				appendEvent(s, "S1", "1", "d3")
				appendEvent(s, "S2", "8", "d4")
				s.SetMaxBytes(2)
				// Up to here, identical to the "purge" case.
				s.SetMaxBytes(6) // make room
				appendEvent(s, "S1", "2", "d5")
				appendEvent(s, "S1", "2", "d6")
			},
			// The other streams remain, because we may add to them.
			"S1 1 first=1 d3; S1 2 first=1 d5 d6; S2 8 first=1",
			6,
		},
	} {
		t.Run(tt.name, func(t *testing.T) {
			s := NewMemoryEventStore(nil)
			tt.actions(s)
			got := s.debugString()
			if got != tt.want {
				t.Errorf("\ngot  %s\nwant %s", got, tt.want)
			}
			if g, w := s.nBytes, tt.wantSize; g != w {
				t.Errorf("got size %d, want %d", g, w)
			}
		})
	}
}

func TestMemoryEventStoreAfter(t *testing.T) {
	ctx := context.Background()
	s := NewMemoryEventStore(nil)
	s.SetMaxBytes(4)
	s.Append(ctx, "S1", "1", []byte("d1"))
	s.Append(ctx, "S1", "1", []byte("d2"))
	s.Append(ctx, "S1", "1", []byte("d3"))
	s.Append(ctx, "S1", "2", []byte("d4")) // will purge d1
	want := "S1 1 first=1 d2 d3; S1 2 first=0 d4"
	if got := s.debugString(); got != want {
		t.Fatalf("got state %q, want %q", got, want)
	}

	for _, tt := range []struct {
		sessionID string
		streamID  string
		index     int
		want      []string
		wantErr   string // if non-empty, error should contain this string
	}{
		{"S1", "1", 0, []string{"d2", "d3"}, ""},
		{"S1", "1", 1, []string{"d3"}, ""},
		{"S1", "1", 2, nil, ""},
		{"S1", "1", 3, nil, ""},
		{"S1", "1", 100, nil, ""}, // past the end of stream "1" (highest is 2)
		{"S1", "2", 0, nil, ""},
		{"S1", "2", 1, nil, ""},
		{"S1", "2", 100, nil, ""},
		{"S1", "3", 0, nil, "unknown stream ID"},
		{"S2", "0", 0, nil, "unknown session ID"},
	} {
		t.Run(fmt.Sprintf("%s-%s-%d", tt.sessionID, tt.streamID, tt.index), func(t *testing.T) {
			var got []string
			for d, err := range s.After(ctx, tt.sessionID, tt.streamID, tt.index) {
				if err != nil {
					if tt.wantErr == "" {
						t.Fatalf("unexpected error %q", err)
					} else if g := err.Error(); !strings.Contains(g, tt.wantErr) {
						t.Fatalf("got error %q, want it to contain %q", g, tt.wantErr)
					} else {
						return
					}
				}
				got = append(got, string(d))
			}
			if tt.wantErr != "" {
				t.Fatalf("expected error containing %q, got nil", tt.wantErr)
			}
			if !slices.Equal(got, tt.want) {
				t.Errorf("got %v, want %v", got, tt.want)
			}
		})
	}
}

func TestMemoryEventStorePurgeEmpty(t *testing.T) {
	ctx := context.Background()
	s := NewMemoryEventStore(nil)
	s.SetMaxBytes(100)

	// Append an empty chunk first.
	if err := s.Append(ctx, "S1", "1", []byte("")); err != nil {
		t.Fatal(err)
	}
	// Append a non-empty chunk.
	if err := s.Append(ctx, "S1", "1", []byte("1234567890")); err != nil {
		t.Fatal(err)
	}
	// Now nBytes = 10, but the first chunk is empty.

	// This should not panic.
	s.SetMaxBytes(5)

	if s.nBytes > 5 {
		t.Errorf("got nBytes %d, want <= 5", s.nBytes)
	}
}

func BenchmarkMemoryEventStore(b *testing.B) {
	// Benchmark with various settings for event store size, number of session,
	// and payload size.
	//
	// Assume a small number of streams per session, which is probably realistic.
	tests := []struct {
		name     string
		limit    int
		sessions int
		datasize int
	}{
		{"1KB", 1024, 1, 16},
		{"1MB", 1024 * 1024, 10, 16},
		{"10MB", 10 * 1024 * 1024, 100, 16},
		{"10MB_big", 10 * 1024 * 1024, 1000, 128},
	}

	for _, test := range tests {
		b.Run(test.name, func(b *testing.B) {
			store := NewMemoryEventStore(nil)
			store.SetMaxBytes(test.limit)
			ctx := context.Background()
			sessionIDs := make([]string, test.sessions)
			streamIDs := make([][3]string, test.sessions)
			for i := range sessionIDs {
				sessionIDs[i] = fmt.Sprint(i)
				for j := range 3 {
					streamIDs[i][j] = rand.Text()
				}
			}
			payload := make([]byte, test.datasize)
			start := time.Now()
			b.ResetTimer()
			for i := range b.N {
				sessionID := sessionIDs[i%len(sessionIDs)]
				streamID := streamIDs[i%len(sessionIDs)][i%3]
				store.Append(ctx, sessionID, streamID, payload)
			}
			b.ReportMetric(float64(test.datasize)*float64(b.N)/time.Since(start).Seconds(), "bytes/s")
		})
	}
}