package mcpgateway

import (
	"context"
	"crypto/rand"
	"io"
	"net/http"
	"strings"
	"sync"
	"testing"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/enterpilot/gomodel/config"
	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/encryption"
)

// newVirtualTestService serves alpha, beta, and gamma upstreams plus the
// given virtual servers.
func newVirtualTestService(t *testing.T, mutateBeta func(*ServerSpec), virtuals ...VirtualServerSpec) (*Service, string) {
	t.Helper()
	service, gatewayURL := newTestService(t, nil,
		testSpec("alpha", newTestUpstream(t, "alpha", addEchoTool("echo")), nil),
		testSpec("beta", newTestUpstream(t, "beta", addEchoTool("search")), mutateBeta),
		testSpec("gamma", newTestUpstream(t, "gamma", addEchoTool("fetch")), nil),
	)
	service.virtualSpecs = make(map[string]VirtualServerSpec, len(virtuals))
	for _, virtual := range virtuals {
		service.virtualSpecs[virtual.Name] = virtual
	}
	return service, gatewayURL
}

func TestVirtualServerServesOnlyMembers(t *testing.T) {
	_, gatewayURL := newVirtualTestService(t, nil,
		VirtualServerSpec{Name: "coding", Servers: []string{"alpha", "beta", "ghost"}})

	session := connectClient(t, gatewayURL+"/mcp/coding", nil)
	assert.Equal(t, []string{"alpha_echo", "beta_search"}, listToolNames(t, session), "members only, namespaced; a missing member is skipped")
	assert.Contains(t, session.InitializeResult().Instructions, `virtual server "coding" aggregating 2 server(s): alpha, beta`)

	for _, name := range []string{"beta_search", "search"} {
		result, err := session.CallTool(context.Background(), &mcp.CallToolParams{Name: name, Arguments: map[string]any{"q": 1}})
		require.NoError(t, err)
		require.False(t, result.IsError)
		text, ok := result.Content[0].(*mcp.TextContent)
		require.True(t, ok)
		assert.Equal(t, `echo:{"q":1}`, text.Text, "namespaced and unique bare names both relay")
	}

	_, err := session.CallTool(context.Background(), &mcp.CallToolParams{Name: "gamma_fetch"})
	require.Error(t, err, "a non-member tool is not callable")
}

func TestVirtualServerNeverWidensAccess(t *testing.T) {
	_, gatewayURL := newVirtualTestService(t, func(spec *ServerSpec) {
		spec.UserPaths = []string{"/team-b"}
		spec.AllowedTools = []string{"other"}
	}, VirtualServerSpec{Name: "coding", Servers: []string{"alpha", "beta"}})

	teamA := connectClient(t, gatewayURL+"/mcp/coding", map[string]string{core.UserPathHeader: "/team-a"})
	assert.Equal(t, []string{"alpha_echo"}, listToolNames(t, teamA), "a member outside the caller's user path stays hidden")

	teamB := connectClient(t, gatewayURL+"/mcp/coding", map[string]string{core.UserPathHeader: "/team-b"})
	assert.Equal(t, []string{"alpha_echo"}, listToolNames(t, teamB), "the member's tool filters still apply")

	narrowed := connectClient(t, gatewayURL+"/mcp/coding", map[string]string{ScopeHeader: "gamma"})
	assert.Empty(t, listToolNames(t, narrowed), "the scope header narrows the view and cannot add a non-member")
	assert.Contains(t, narrowed.InitializeResult().Instructions, "None of its member servers is available for this API key.")
}

func TestVirtualServerToolDiscoveryDefault(t *testing.T) {
	service, gatewayURL := newVirtualTestService(t, nil,
		VirtualServerSpec{Name: "search", Servers: []string{"alpha"}, ToolDiscovery: config.MCPToolDiscoverySearch},
		VirtualServerSpec{Name: "listed", Servers: []string{"alpha"}, ToolDiscovery: config.MCPToolDiscoveryOff},
		VirtualServerSpec{Name: "inherit", Servers: []string{"alpha"}},
	)
	service.searchDiscovery = true

	tests := []struct {
		endpoint string
		header   string
		want     []string
	}{
		{endpoint: "search", want: []string{callToolName, searchToolsName}},
		{endpoint: "listed", want: []string{"alpha_echo"}},
		{endpoint: "inherit", want: []string{callToolName, searchToolsName}},
		{endpoint: "search", header: config.MCPToolDiscoveryOff, want: []string{"alpha_echo"}},
	}
	for _, tt := range tests {
		t.Run(tt.endpoint+"/"+tt.header, func(t *testing.T) {
			headers := map[string]string{}
			if tt.header != "" {
				headers[ToolDiscoveryHeader] = tt.header
			}
			session := connectClient(t, gatewayURL+"/mcp/"+tt.endpoint, headers)
			assert.Equal(t, tt.want, listToolNames(t, session))
		})
	}
}

func TestVirtualServerYieldsToServerWithSameName(t *testing.T) {
	// Started before the service so cleanup closes the service first.
	codingURL := newTestUpstream(t, "coding", addEchoTool("lint"))
	service, gatewayURL := newVirtualTestService(t, nil,
		VirtualServerSpec{Name: "coding", Servers: []string{"alpha"}})

	// An admin-managed server claims the name after the virtual server exists.
	specs := make([]ServerSpec, 0, 4)
	for _, view := range service.Views() {
		specs = append(specs, view.Spec)
	}
	specs = append(specs, testSpec("coding", codingURL, func(spec *ServerSpec) {
		spec.Managed = false
	}))
	service.manager.Apply(specs)
	waitForConnected(t, service, len(specs))

	session := connectClient(t, gatewayURL+"/mcp/coding", nil)
	assert.Equal(t, []string{"lint"}, listToolNames(t, session), "the real server keeps its endpoint")

	views := service.VirtualViews()
	require.Len(t, views, 1)
	assert.Equal(t, `virtual MCP server "coding" is not served: MCP server "coding" uses the same name; rename one of them`, views[0].Conflict)
}

func TestVirtualViewsReportMembersAndDiscovery(t *testing.T) {
	service, _ := newVirtualTestService(t, nil,
		VirtualServerSpec{Name: "coding", Servers: []string{"alpha", "ghost"}, ToolDiscovery: config.MCPToolDiscoverySearch},
		VirtualServerSpec{Name: "all", Servers: []string{"alpha", "beta", "gamma"}},
	)

	views := service.VirtualViews()
	require.Len(t, views, 2)
	assert.Equal(t, "all", views[0].Spec.Name, "sorted by name")
	assert.Equal(t, config.MCPToolDiscoveryOff, views[0].ToolDiscovery)
	assert.Empty(t, views[0].MissingServers)
	assert.Empty(t, views[0].Conflict)
	assert.Equal(t, config.MCPToolDiscoverySearch, views[1].ToolDiscovery)
	assert.Equal(t, []string{"ghost"}, views[1].MissingServers)
}

func TestUnknownEndpointErrorNamesBothKinds(t *testing.T) {
	_, gatewayURL := newVirtualTestService(t, nil)

	resp, err := http.Post(gatewayURL+"/mcp/ghost", "application/json", strings.NewReader(`{}`))
	require.NoError(t, err)
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Equal(t, http.StatusNotFound, resp.StatusCode)
	assert.Contains(t, string(body), "unknown MCP server or virtual server: ghost")
}

func TestSessionBindingRejectsSwitchBetweenVirtualAndPinned(t *testing.T) {
	_, gatewayURL := newVirtualTestService(t, nil,
		VirtualServerSpec{Name: "coding", Servers: []string{"alpha"}})

	listBody := `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`
	virtualSession := initializeRawSession(t, gatewayURL+"/mcp/coding", nil)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, gatewayURL+"/mcp/alpha", listBody, map[string]string{"Mcp-Session-Id": virtualSession}))
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, gatewayURL+"/mcp", listBody, map[string]string{"Mcp-Session-Id": virtualSession}))

	pinnedSession := initializeRawSession(t, gatewayURL+"/mcp/alpha", nil)
	assert.Equal(t, http.StatusNotFound, rawMCPStatus(t, gatewayURL+"/mcp/coding", listBody, map[string]string{"Mcp-Session-Id": pinnedSession}))
}

// memoryStore is a minimal in-process Store for admin write paths.
type memoryStore struct{ rows map[string]ManagedServer }

func (m *memoryStore) List(context.Context) ([]ManagedServer, error) {
	rows := make([]ManagedServer, 0, len(m.rows))
	for _, row := range m.rows {
		rows = append(rows, row)
	}
	return rows, nil
}

func (m *memoryStore) Get(_ context.Context, name string) (*ManagedServer, error) {
	row, ok := m.rows[name]
	if !ok {
		return nil, ErrNotFound
	}
	return &row, nil
}

func (m *memoryStore) Upsert(_ context.Context, server ManagedServer) error {
	m.rows[server.Name] = server
	return nil
}

func (m *memoryStore) Update(_ context.Context, server ManagedServer) error {
	if _, ok := m.rows[server.Name]; !ok {
		return ErrNotFound
	}
	m.rows[server.Name] = server
	return nil
}

func (m *memoryStore) Delete(_ context.Context, name string) error {
	delete(m.rows, name)
	return nil
}

func (m *memoryStore) Close() error { return nil }

// memoryVirtualStore is a minimal in-process VirtualStore. It is locked
// because the background refresh reads it while tests write.
type memoryVirtualStore struct {
	mu   sync.Mutex
	rows map[string]ManagedVirtualServer
}

func (m *memoryVirtualStore) put(virtual ManagedVirtualServer) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.rows[virtual.Name] = virtual
}

func (m *memoryVirtualStore) ListVirtual(context.Context) ([]ManagedVirtualServer, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	rows := make([]ManagedVirtualServer, 0, len(m.rows))
	for _, row := range m.rows {
		rows = append(rows, row)
	}
	return rows, nil
}

func (m *memoryVirtualStore) GetVirtual(_ context.Context, name string) (*ManagedVirtualServer, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	row, ok := m.rows[name]
	if !ok {
		return nil, ErrNotFound
	}
	return &row, nil
}

// UpsertVirtual stamps a generation on insert and keeps the row's own on an
// edit, like the real stores.
func (m *memoryVirtualStore) UpsertVirtual(_ context.Context, virtual ManagedVirtualServer) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if existing, ok := m.rows[virtual.Name]; ok {
		virtual.Generation = existing.Generation
	} else if virtual.Generation == "" {
		virtual.Generation = rand.Text()
	}
	m.rows[virtual.Name] = virtual
	return nil
}

func (m *memoryVirtualStore) UpdateVirtual(_ context.Context, virtual ManagedVirtualServer) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.rows[virtual.Name]; !ok {
		return ErrNotFound
	}
	m.rows[virtual.Name] = virtual
	return nil
}

func (m *memoryVirtualStore) DeleteVirtual(_ context.Context, name string) (*ManagedVirtualServer, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	row, ok := m.rows[name]
	if !ok {
		return nil, ErrNotFound
	}
	delete(m.rows, name)
	return &row, nil
}

type unconfirmedStore struct{ *memoryStore }

func (s unconfirmedStore) Upsert(ctx context.Context, server ManagedServer) error {
	if err := s.memoryStore.Upsert(ctx, server); err != nil {
		return err
	}
	return encryption.ErrSealUnconfirmed
}

func TestUpsertAppliesASaveWhoseDataKeyIsUnconfirmed(t *testing.T) {
	store := unconfirmedStore{&memoryStore{rows: map[string]ManagedServer{}}}
	service, err := NewService(context.Background(), Options{Store: store})
	require.NoError(t, err)
	t.Cleanup(service.Close)

	err = service.Upsert(context.Background(), ManagedServer{Name: "fresh", URL: "https://example.com/mcp", Transport: config.MCPTransportHTTP})
	require.ErrorIs(t, err, encryption.ErrSealUnconfirmed, "the caller learns the save is not confirmed")
	_, applied := service.manager.get("fresh")
	assert.True(t, applied, "the saved row is applied like any other save")
}