package mcpserver

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"testing"
	"time"

	"email-mcp/internal/imapclient"
	"email-mcp/internal/secretstore"
	"email-mcp/internal/secretstore/kwallet"
)

type storeStub struct {
	credential secretstore.Credential
	loadErr    error
	loadCalls  int
	loadKey    string
}

func (s *storeStub) Save(context.Context, string, secretstore.Credential) error {
	return nil
}

func (s *storeStub) Load(_ context.Context, key string) (secretstore.Credential, error) {
	s.loadCalls++
	s.loadKey = key
	if s.loadErr != nil {
		return secretstore.Credential{}, s.loadErr
	}
	return s.credential, nil
}

type serviceStub struct {
	listMailboxes func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error)
	listMessages  func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error)
	getMessage    func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error)
}

func (s serviceStub) ListMailboxes(ctx context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) {
	return s.listMailboxes(ctx, cred)
}

func (s serviceStub) ListMessages(ctx context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
	return s.listMessages(ctx, cred, mailbox, limit)
}

func (s serviceStub) GetMessage(ctx context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) {
	return s.getMessage(ctx, cred, mailbox, uid)
}

func TestServerListMailboxesLoadsCredentialAndDelegates(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	server := New(store, serviceStub{
		listMailboxes: func(_ context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) {
			if cred.Host != "imap.example.com" || cred.Username != "alice" || cred.Password != "secret" {
				t.Fatalf("unexpected credential: %#v", cred)
			}
			return []imapclient.Mailbox{{Name: "INBOX"}}, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	})

	result, err := server.ListMailboxes(context.Background())
	if err != nil {
		t.Fatalf("ListMailboxes returned error: %v", err)
	}
	if store.loadCalls != 1 {
		t.Fatalf("expected credential to be loaded once, got %d", store.loadCalls)
	}
	if store.loadKey != secretstore.DefaultAccountKey {
		t.Fatalf("expected load key %q, got %q", secretstore.DefaultAccountKey, store.loadKey)
	}
	if len(result) != 1 || result[0].Name != "INBOX" {
		t.Fatalf("unexpected result: %#v", result)
	}
}

func TestServerListMessagesLoadsCredentialAndDelegates(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	server := New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
			if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != 5 {
				t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit)
			}
			return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	})

	result, err := server.ListMessages(context.Background(), "INBOX", 5)
	if err != nil {
		t.Fatalf("ListMessages returned error: %v", err)
	}
	if len(result) != 1 || result[0].UID != 42 {
		t.Fatalf("unexpected result: %#v", result)
	}
}

func TestServerGetMessageUsesUIDContract(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	server := New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(_ context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) {
			if cred.Host != "imap.example.com" || mailbox != "INBOX" || uid != 42 {
				t.Fatalf("unexpected call: cred=%#v mailbox=%q uid=%d", cred, mailbox, uid)
			}
			return imapclient.Message{
				UID:     42,
				Mailbox: "INBOX",
				Body:    "body",
			}, nil
		},
	})

	message, err := server.GetMessage(context.Background(), "INBOX", 42)
	if err != nil {
		t.Fatalf("GetMessage returned error: %v", err)
	}
	if message.UID != 42 || message.Mailbox != "INBOX" {
		t.Fatalf("unexpected message: %#v", message)
	}
}

func TestRunnerRunWritesToolManifestAndHandlesRequests(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := bytes.NewBufferString("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test-client\",\"version\":\"1.0.0\"}}}\n{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/list\",\"params\":{}}\n{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"tools/call\",\"params\":{\"name\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\",\"limit\":5}}}\n")
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
			if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != 5 {
				t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit)
			}
			return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}
	if store.loadCalls != 1 {
		t.Fatalf("expected credential preload once, got %d", store.loadCalls)
	}

	decoder := json.NewDecoder(output)

	var initializeResponse struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Result  struct {
			ProtocolVersion string         `json:"protocolVersion"`
			Capabilities    map[string]any `json:"capabilities"`
			ServerInfo      struct {
				Name    string `json:"name"`
				Version string `json:"version"`
			} `json:"serverInfo"`
		} `json:"result"`
	}
	if err := decoder.Decode(&initializeResponse); err != nil {
		t.Fatalf("failed to decode initialize response: %v", err)
	}
	if initializeResponse.JSONRPC != "2.0" {
		t.Fatalf("expected jsonrpc 2.0, got %#v", initializeResponse)
	}
	if initializeResponse.ID != 1 {
		t.Fatalf("expected initialize response id 1, got %#v", initializeResponse)
	}
	if initializeResponse.Result.ProtocolVersion != "2025-03-26" {
		t.Fatalf("expected negotiated protocol version, got %#v", initializeResponse.Result)
	}
	if _, ok := initializeResponse.Result.Capabilities["tools"]; !ok {
		t.Fatalf("expected tools capability, got %#v", initializeResponse.Result.Capabilities)
	}
	if initializeResponse.Result.ServerInfo.Name == "" || initializeResponse.Result.ServerInfo.Version == "" {
		t.Fatalf("expected server info, got %#v", initializeResponse.Result.ServerInfo)
	}

	var listResponse struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Result  struct {
			Tools []map[string]any `json:"tools"`
		} `json:"result"`
	}
	if err := decoder.Decode(&listResponse); err != nil {
		t.Fatalf("failed to decode tools/list response: %v", err)
	}
	if listResponse.JSONRPC != "2.0" || listResponse.ID != 2 {
		t.Fatalf("unexpected tools/list response envelope: %#v", listResponse)
	}
	if len(listResponse.Result.Tools) != 3 {
		t.Fatalf("expected 3 tools, got %#v", listResponse.Result.Tools)
	}
	if listResponse.Result.Tools[0]["name"] != "list_mailboxes" || listResponse.Result.Tools[1]["name"] != "list_messages" || listResponse.Result.Tools[2]["name"] != "get_message" {
		t.Fatalf("unexpected tool manifest: %#v", listResponse.Result.Tools)
	}
	if _, ok := listResponse.Result.Tools[1]["inputSchema"]; !ok {
		t.Fatalf("expected inputSchema field in tools/list response, got %#v", listResponse.Result.Tools[1])
	}

	var response struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Result  struct {
			Content []struct {
				Type string `json:"type"`
				Text string `json:"text"`
			} `json:"content"`
			IsError bool `json:"isError"`
		} `json:"result"`
	}
	if err := decoder.Decode(&response); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if response.JSONRPC != "2.0" || response.ID != 3 {
		t.Fatalf("unexpected tools/call response envelope: %#v", response)
	}
	if response.Result.IsError {
		t.Fatalf("expected successful tools/call result, got %#v", response.Result)
	}
	if len(response.Result.Content) != 1 || response.Result.Content[0].Type != "text" {
		t.Fatalf("unexpected tools/call content: %#v", response.Result.Content)
	}

	var messages []imapclient.MessageSummary
	if err := json.Unmarshal([]byte(response.Result.Content[0].Text), &messages); err != nil {
		t.Fatalf("failed to decode tools/call text payload: %v", err)
	}
	if len(messages) != 1 || messages[0].UID != 42 {
		t.Fatalf("unexpected response: %#v", messages)
	}
}

func TestRunnerRunAcceptsClaudeCodeProtocolVersion(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := bytes.NewBufferString("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{},\"clientInfo\":{\"name\":\"claude-code\",\"version\":\"1.0.0\"}}}\n")
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}

	var response struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Result  struct {
			ProtocolVersion string `json:"protocolVersion"`
		} `json:"result"`
		Error *struct {
			Code    int    `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.NewDecoder(output).Decode(&response); err != nil {
		t.Fatalf("failed to decode initialize response: %v", err)
	}
	if response.Error != nil {
		t.Fatalf("expected initialize to succeed, got error %#v", response.Error)
	}
	if response.JSONRPC != "2.0" || response.ID != 1 {
		t.Fatalf("unexpected response envelope: %#v", response)
	}
	if response.Result.ProtocolVersion != "2024-11-05" {
		t.Fatalf("expected negotiated protocol version 2024-11-05, got %#v", response.Result)
	}
}

func TestRunnerRunReturnsFriendlyMissingCredentialError(t *testing.T) {
	store := &storeStub{
		loadErr: kwallet.ErrCredentialNotFound,
	}
	input := bytes.NewBufferString(
		"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test\",\"version\":\"1.0.0\"}}}\n" +
			"{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n" +
			"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"list_mailboxes\"}}\n",
	)
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}

	decoder := json.NewDecoder(output)

	// Skip initialize response
	var initResp json.RawMessage
	if err := decoder.Decode(&initResp); err != nil {
		t.Fatalf("failed to decode initialize response: %v", err)
	}

	// Check tool call response contains credential error
	var toolResp struct {
		ID     int `json:"id"`
		Result struct {
			Content []struct {
				Text string `json:"text"`
			} `json:"content"`
			IsError bool `json:"isError"`
		} `json:"result"`
	}
	if err := decoder.Decode(&toolResp); err != nil {
		t.Fatalf("failed to decode tool call response: %v", err)
	}
	if !toolResp.Result.IsError {
		t.Fatal("expected isError true for missing credentials")
	}
	if len(toolResp.Result.Content) == 0 || toolResp.Result.Content[0].Text != ErrCredentialsNotConfigured.Error() {
		t.Fatalf("expected credential error message, got %#v", toolResp.Result)
	}
}

func TestRunnerRunReturnsFriendlyMissingCredentialErrorWhenStoreAlreadyTranslatedIt(t *testing.T) {
	store := &storeStub{
		loadErr: ErrCredentialsNotConfigured,
	}
	input := bytes.NewBufferString(
		"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test\",\"version\":\"1.0.0\"}}}\n" +
			"{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n" +
			"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"list_mailboxes\"}}\n",
	)
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}

	decoder := json.NewDecoder(output)

	// Skip initialize response
	var initResp json.RawMessage
	if err := decoder.Decode(&initResp); err != nil {
		t.Fatalf("failed to decode initialize response: %v", err)
	}

	// Check tool call response contains credential error
	var toolResp struct {
		ID     int `json:"id"`
		Result struct {
			Content []struct {
				Text string `json:"text"`
			} `json:"content"`
			IsError bool `json:"isError"`
		} `json:"result"`
	}
	if err := decoder.Decode(&toolResp); err != nil {
		t.Fatalf("failed to decode tool call response: %v", err)
	}
	if !toolResp.Result.IsError {
		t.Fatal("expected isError true for missing credentials")
	}
	if len(toolResp.Result.Content) == 0 || toolResp.Result.Content[0].Text != ErrCredentialsNotConfigured.Error() {
		t.Fatalf("expected credential error message, got %#v", toolResp.Result)
	}
}

func TestServerToolsAdvertiseValidatedArgumentContracts(t *testing.T) {
	tools := New(&storeStub{}, serviceStub{}).Tools()
	if len(tools) != 3 {
		t.Fatalf("expected 3 tools, got %d", len(tools))
	}

	listMessages := tools[1]
	if listMessages.Name != "list_messages" {
		t.Fatalf("unexpected tool ordering: %#v", tools)
	}
	if got := listMessages.InputSchema["type"]; got != "object" {
		t.Fatalf("expected object schema, got %#v", got)
	}

	listProps, ok := listMessages.InputSchema["properties"].(map[string]any)
	if !ok {
		t.Fatalf("expected properties map, got %#v", listMessages.InputSchema["properties"])
	}
	limitSchema, ok := listProps["limit"].(map[string]any)
	if !ok {
		t.Fatalf("expected limit schema, got %#v", listProps["limit"])
	}
	mailboxSchema, ok := listProps["mailbox"].(map[string]any)
	if !ok {
		t.Fatalf("expected mailbox schema, got %#v", listProps["mailbox"])
	}
	if got := mailboxSchema["pattern"]; got != "\\S" {
		t.Fatalf("expected mailbox pattern %q, got %#v", "\\S", got)
	}
	if got := limitSchema["default"]; got != float64(imapclient.DefaultListMessagesLimit) && got != imapclient.DefaultListMessagesLimit {
		t.Fatalf("expected limit default %d, got %#v", imapclient.DefaultListMessagesLimit, got)
	}
	if got := limitSchema["minimum"]; got != float64(1) && got != 1 {
		t.Fatalf("expected limit minimum 1, got %#v", got)
	}
	if got := limitSchema["maximum"]; got != float64(imapclient.MaxListMessagesLimit) && got != imapclient.MaxListMessagesLimit {
		t.Fatalf("expected limit maximum %d, got %#v", imapclient.MaxListMessagesLimit, got)
	}

	getMessage := tools[2]
	if getMessage.Name != "get_message" {
		t.Fatalf("unexpected tool ordering: %#v", tools)
	}
	getProps, ok := getMessage.InputSchema["properties"].(map[string]any)
	if !ok {
		t.Fatalf("expected get_message properties map, got %#v", getMessage.InputSchema["properties"])
	}
	getMailboxSchema, ok := getProps["mailbox"].(map[string]any)
	if !ok {
		t.Fatalf("expected get_message mailbox schema, got %#v", getProps["mailbox"])
	}
	if got := getMailboxSchema["pattern"]; got != "\\S" {
		t.Fatalf("expected mailbox pattern %q, got %#v", "\\S", got)
	}
	uidSchema, ok := getProps["uid"].(map[string]any)
	if !ok {
		t.Fatalf("expected uid schema, got %#v", getProps["uid"])
	}
	if got := uidSchema["minimum"]; got != float64(1) && got != 1 {
		t.Fatalf("expected uid minimum 1, got %#v", got)
	}
}

func TestRunnerRunReturnsValidationErrorsForInvalidRequests(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := bytes.NewBufferString("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test-client\",\"version\":\"1.0.0\"}}}\n{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\",\"limit\":0}}}\n{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"tools/call\",\"params\":{\"name\":\"get_message\",\"arguments\":{\"mailbox\":\"INBOX\"}}}\n")
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}

	decoder := json.NewDecoder(output)
	if err := decoder.Decode(&struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
	}{}); err != nil {
		t.Fatalf("failed to decode initialize response: %v", err)
	}

	var firstResponse struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Error   struct {
			Code    int    `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := decoder.Decode(&firstResponse); err != nil {
		t.Fatalf("failed to decode first error response: %v", err)
	}
	if firstResponse.JSONRPC != "2.0" || firstResponse.ID != 2 {
		t.Fatalf("unexpected first error envelope: %#v", firstResponse)
	}
	if firstResponse.Error.Code != -32602 {
		t.Fatalf("expected invalid params code, got %#v", firstResponse)
	}
	if firstResponse.Error.Message != "limit must be between 1 and 50" {
		t.Fatalf("unexpected first error: %#v", firstResponse)
	}

	var secondResponse struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Error   struct {
			Code    int    `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := decoder.Decode(&secondResponse); err != nil {
		t.Fatalf("failed to decode second error response: %v", err)
	}
	if secondResponse.JSONRPC != "2.0" || secondResponse.ID != 3 {
		t.Fatalf("unexpected second error envelope: %#v", secondResponse)
	}
	if secondResponse.Error.Code != -32602 {
		t.Fatalf("expected invalid params code, got %#v", secondResponse)
	}
	if secondResponse.Error.Message != "uid is required" {
		t.Fatalf("unexpected second error: %#v", secondResponse)
	}
}

func TestRunnerRunRejectsWhitespaceOnlyMailboxValues(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := bytes.NewBufferString("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test-client\",\"version\":\"1.0.0\"}}}\n{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"list_messages\",\"arguments\":{\"mailbox\":\"   \"}}}\n")
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}

	decoder := json.NewDecoder(output)
	if err := decoder.Decode(&struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
	}{}); err != nil {
		t.Fatalf("failed to decode initialize response: %v", err)
	}

	var response struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Error   struct {
			Code    int    `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := decoder.Decode(&response); err != nil {
		t.Fatalf("failed to decode error response: %v", err)
	}
	if response.JSONRPC != "2.0" || response.ID != 2 {
		t.Fatalf("unexpected error envelope: %#v", response)
	}
	if response.Error.Code != -32602 {
		t.Fatalf("expected invalid params code, got %#v", response)
	}
	if response.Error.Message != "mailbox is required" {
		t.Fatalf("unexpected error: %#v", response)
	}
}

func TestRunnerRunAppliesDefaultLimitWhenOmitted(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := bytes.NewBufferString("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{},\"clientInfo\":{\"name\":\"test-client\",\"version\":\"1.0.0\"}}}\n{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\",\"params\":{\"name\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\"}}}\n")
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
			if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != imapclient.DefaultListMessagesLimit {
				t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit)
			}
			return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}
}

func TestRunnerRunRejectsRequestsBeforeInitialize(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := bytes.NewBufferString("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/list\",\"params\":{}}\n")
	output := &bytes.Buffer{}
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, output, &bytes.Buffer{})

	if err := runner.Run(context.Background()); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}

	var response struct {
		JSONRPC string `json:"jsonrpc"`
		ID      int    `json:"id"`
		Error   struct {
			Code    int    `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.NewDecoder(output).Decode(&response); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if response.JSONRPC != "2.0" || response.ID != 1 {
		t.Fatalf("unexpected response envelope: %#v", response)
	}
	if response.Error.Code == 0 || response.Error.Message == "" {
		t.Fatalf("expected protocol error before initialization, got %#v", response)
	}
}

func TestRunnerRunStopsWhenContextCanceledWhileWaitingForInput(t *testing.T) {
	store := &storeStub{
		credential: secretstore.Credential{
			Host:     "imap.example.com",
			Username: "alice",
			Password: "secret",
		},
	}
	input := newBlockingReadCloser()
	runner := NewRunner(New(store, serviceStub{
		listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
			t.Fatal("ListMailboxes should not be called")
			return nil, nil
		},
		listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
			t.Fatal("ListMessages should not be called")
			return nil, nil
		},
		getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
			t.Fatal("GetMessage should not be called")
			return imapclient.Message{}, nil
		},
	}), input, &bytes.Buffer{}, &bytes.Buffer{})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	done := make(chan error, 1)
	go func() {
		done <- runner.Run(ctx)
	}()

	select {
	case <-input.started:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("runner never started reading input")
	}

	cancel()

	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("expected context cancellation, got %v", err)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("runner did not stop after context cancellation")
	}

	if !input.closed {
		t.Fatal("expected runner to close input reader on cancellation")
	}
}

type blockingReadCloser struct {
	started chan struct{}
	closed  bool
	done    chan struct{}
}

func newBlockingReadCloser() *blockingReadCloser {
	return &blockingReadCloser{
		started: make(chan struct{}),
		done:    make(chan struct{}),
	}
}

func (r *blockingReadCloser) Read(_ []byte) (int, error) {
	select {
	case <-r.started:
	default:
		close(r.started)
	}
	<-r.done
	return 0, io.EOF
}

func (r *blockingReadCloser) Close() error {
	if !r.closed {
		r.closed = true
		close(r.done)
	}
	return nil
}