package mcpgateway

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"slices"
	"testing"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"

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

	"github.com/enterpilot/gomodel/config"
)

// newVirtualAdminTestService serves alpha and beta upstreams from config,
// with in-memory stores for servers and virtual servers.
func newVirtualAdminTestService(t *testing.T, configVirtuals map[string]VirtualServerSpec, stored ...ManagedVirtualServer) (*Service, *memoryVirtualStore, string) {
	t.Helper()
	virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{}}
	for _, row := range stored {
		virtualStore.rows[row.Name] = row
	}
	service, gatewayURL := newTestServiceWithOptions(t, Options{
		ConfigServers: map[string]ServerSpec{
			"alpha": testSpec("alpha", newTestUpstream(t, "alpha", addEchoTool("echo")), nil),
			"beta":  testSpec("beta", newTestUpstream(t, "beta", addEchoTool("search")), nil),
		},
		Store:          &memoryStore{rows: map[string]ManagedServer{}},
		VirtualServers: configVirtuals,
		VirtualStore:   virtualStore,
	})
	return service, virtualStore, gatewayURL
}

func TestUpsertVirtualServesImmediately(t *testing.T) {
	service, virtualStore, gatewayURL := newVirtualAdminTestService(t, nil)

	view, err := service.UpsertVirtual(context.Background(), ManagedVirtualServer{
		Name:          " Coding ",
		Description:   "code tools",
		Servers:       []string{"beta", "alpha"},
		ToolDiscovery: config.MCPToolDiscoverySearch,
	})
	require.NoError(t, err)
	require.Contains(t, virtualStore.rows, "coding", "the name is normalized before it is stored")
	assert.Equal(t, "coding", view.Spec.Name, "the saved definition is returned")
	assert.Equal(t, config.MCPToolDiscoverySearch, view.ToolDiscovery)

	views := service.VirtualViews()
	require.Len(t, views, 1)
	assert.False(t, views[0].Spec.Managed)
	assert.Equal(t, []string{"beta", "alpha"}, views[0].Spec.Servers)

	session := connectClient(t, gatewayURL+"/mcp/coding", map[string]string{ToolDiscoveryHeader: config.MCPToolDiscoveryOff})
	assert.Equal(t, []string{"alpha_echo", "beta_search"}, listToolNames(t, session))
}

func TestUpsertVirtualRejections(t *testing.T) {
	tests := []struct {
		name    string
		virtual ManagedVirtualServer
		wantErr string
	}{
		{
			name:    "unknown member",
			virtual: ManagedVirtualServer{Name: "coding", Servers: []string{"alpha", "ghost"}},
			wantErr: `member "ghost" matches no MCP server`,
		},
		{
			name:    "nested virtual server",
			virtual: ManagedVirtualServer{Name: "all", Servers: []string{"research"}},
			wantErr: `member "research" is a virtual server; virtual servers can only include MCP servers`,
		},
		{
			name:    "server slug",
			virtual: ManagedVirtualServer{Name: "alpha", Servers: []string{"beta"}},
			wantErr: `name "alpha" is used by MCP server "alpha"; choose another name`,
		},
		{
			name:    "config-declared name",
			virtual: ManagedVirtualServer{Name: "research", Servers: []string{"beta"}},
			wantErr: `virtual MCP server "research" is managed by config/env and is read-only`,
		},
		{
			name:    "invalid name",
			virtual: ManagedVirtualServer{Name: "my tools", Servers: []string{"alpha"}},
			wantErr: `virtual server name "my tools" must match`,
		},
		{
			name:    "no members",
			virtual: ManagedVirtualServer{Name: "coding"},
			wantErr: "servers is required",
		},
		{
			name:    "invalid tool discovery",
			virtual: ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}, ToolDiscovery: "semantic"},
			wantErr: `tool_discovery must be "off" or "search"`,
		},
		{
			// Admin input is literal: a reference is checked, not stored.
			name:    "secret reference in tool discovery",
			virtual: ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}, ToolDiscovery: "${env:MODE}"},
			wantErr: `tool_discovery must be "off" or "search", got "${env:MODE}"`,
		},
		{
			name:    "secret reference as a member",
			virtual: ManagedVirtualServer{Name: "coding", Servers: []string{"${env:MEMBER}"}},
			wantErr: `member "${env:MEMBER}"`,
		},
	}
	service, virtualStore, _ := newVirtualAdminTestService(t, map[string]VirtualServerSpec{
		"research": {Name: "research", Servers: []string{"beta"}, Managed: true},
	})
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := service.UpsertVirtual(context.Background(), tt.virtual)
			require.ErrorIs(t, err, ErrInvalidVirtualServer)
			assert.Contains(t, err.Error(), tt.wantErr)
			assert.Empty(t, virtualStore.rows)
		})
	}
}

func TestUpsertVirtualKeepsRowStoredBeforeServerEditable(t *testing.T) {
	// The row predates the "alpha" server, so it is not served but stays editable.
	service, virtualStore, _ := newVirtualAdminTestService(t, nil,
		ManagedVirtualServer{Name: "alpha", Servers: []string{"beta"}})
	require.NotEmpty(t, service.VirtualViews()[0].Conflict)

	_, err := service.UpsertVirtual(context.Background(), ManagedVirtualServer{Name: "alpha", Description: "renamed soon", Servers: []string{"beta"}})
	require.NoError(t, err)
	assert.Equal(t, "renamed soon", virtualStore.rows["alpha"].Description)

	// If the row is deleted while the edit is in flight, the edit must not
	// recreate it under the server's slug.
	delete(virtualStore.rows, "alpha")
	_, err = service.UpsertVirtual(context.Background(), ManagedVirtualServer{Name: "alpha", Servers: []string{"beta"}})
	require.ErrorIs(t, err, ErrInvalidVirtualServer)
	assert.Empty(t, virtualStore.rows)
}

func TestConfigVirtualServerShadowsStoredRow(t *testing.T) {
	service, _, _ := newVirtualAdminTestService(t,
		map[string]VirtualServerSpec{"coding": {Name: "coding", Servers: []string{"alpha"}, Managed: true}},
		ManagedVirtualServer{Name: "coding", Servers: []string{"beta"}},
	)

	views := service.VirtualViews()
	require.Len(t, views, 1)
	assert.True(t, views[0].Spec.Managed)
	assert.Equal(t, []string{"alpha"}, views[0].Spec.Servers)
	assert.True(t, service.IsManagedVirtual("coding"))
}

func TestDeleteVirtualEndsItsEndpoint(t *testing.T) {
	service, virtualStore, gatewayURL := newVirtualAdminTestService(t,
		map[string]VirtualServerSpec{"research": {Name: "research", Servers: []string{"beta"}, Managed: true}},
		ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}},
	)
	sessionID := initializeRawSession(t, gatewayURL+"/mcp/coding", nil)
	service.bindMu.Lock()
	host := service.bindings[sessionID].server
	service.bindMu.Unlock()
	require.NotNil(t, host)
	require.NotEmpty(t, slices.Collect(host.Sessions()))

	require.NoError(t, service.DeleteVirtual(context.Background(), " Coding "), "names are normalized like on save")
	assert.Empty(t, virtualStore.rows)
	assert.False(t, service.IsVirtual("coding"))

	// The session is closed and unbound now, not at the idle timeout.
	service.bindMu.Lock()
	_, bound := service.bindings[sessionID]
	service.bindMu.Unlock()
	assert.False(t, bound)
	assert.Eventually(t, func() bool { return len(slices.Collect(host.Sessions())) == 0 }, 5*time.Second, 10*time.Millisecond)

	listBody := `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`
	status := rawMCPStatus(t, gatewayURL+"/mcp/coding", listBody, map[string]string{"Mcp-Session-Id": sessionID})
	assert.Equal(t, http.StatusNotFound, status, "a session on a deleted virtual server ends")

	require.ErrorIs(t, service.DeleteVirtual(context.Background(), "coding"), ErrNotFound)
	err := service.DeleteVirtual(context.Background(), "Research")
	require.ErrorIs(t, err, ErrInvalidVirtualServer, "config-declared virtual servers are read-only, whatever the case")
}

func TestUpsertServerRejectsStoredVirtualName(t *testing.T) {
	service, _, _ := newVirtualAdminTestService(t, nil,
		ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})

	err := service.Upsert(context.Background(), ManagedServer{Name: "coding", URL: "https://example.com/mcp", Transport: config.MCPTransportHTTP})
	require.ErrorIs(t, err, ErrVirtualNameTaken)
	assert.Equal(t, `slug "coding" is used by virtual MCP server "coding"; choose another slug`, err.Error())
}

func TestVirtualMemberRemovalReachesOpenSessions(t *testing.T) {
	service, _, gatewayURL := newVirtualAdminTestService(t, nil,
		ManagedVirtualServer{Name: "coding", Servers: []string{"alpha", "beta"}})
	session := connectClient(t, gatewayURL+"/mcp/coding", nil)
	require.Equal(t, []string{"alpha_echo", "beta_search"}, listToolNames(t, session))

	_, err := service.UpsertVirtual(context.Background(), ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
	require.NoError(t, err)

	_, err = session.CallTool(context.Background(), &mcp.CallToolParams{Name: "beta_search"})
	require.Error(t, err, "a removed member must not stay callable from an open session")
	result, err := session.CallTool(context.Background(), &mcp.CallToolParams{Name: "alpha_echo"})
	require.NoError(t, err)
	assert.False(t, result.IsError, "remaining members keep working")
}

func TestBackgroundRefreshPicksUpOtherInstanceWrites(t *testing.T) {
	virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{}}
	service, _ := newTestServiceWithOptions(t, Options{
		ConfigServers:   map[string]ServerSpec{"alpha": testSpec("alpha", newTestUpstream(t, "alpha", addEchoTool("echo")), nil)},
		Store:           &memoryStore{rows: map[string]ManagedServer{}},
		VirtualStore:    virtualStore,
		RefreshInterval: 20 * time.Millisecond,
	})

	// Another gateway instance sharing the database saves a virtual server.
	virtualStore.put(ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
	require.Eventually(t, func() bool { return service.IsVirtual("coding") }, 5*time.Second, 10*time.Millisecond)
}

// gatedStore blocks List while armed, standing in for a slow store read that
// a reload is in the middle of when Close runs. honorCtx makes the read
// return on cancellation, like a real database driver.
type gatedStore struct {
	*memoryStore
	gate     chan struct{}
	entered  chan struct{}
	honorCtx bool
}

func (g *gatedStore) List(ctx context.Context) ([]ManagedServer, error) {
	if g.gate != nil {
		g.entered <- struct{}{}
		if g.honorCtx {
			select {
			case <-g.gate:
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		} else {
			<-g.gate
		}
	}
	return g.memoryStore.List(ctx)
}

func TestCloseWaitsForInFlightReloadAndStopsLaterOnes(t *testing.T) {
	lateURL := newTestUpstream(t, "late", addEchoTool("echo"))
	store := &gatedStore{memoryStore: &memoryStore{rows: map[string]ManagedServer{}}}
	service, err := NewService(context.Background(), Options{Store: store})
	require.NoError(t, err)
	t.Cleanup(service.Close)

	// A server saved on another instance, read by a refresh that is slow.
	store.rows["late"] = ManagedServer{Name: "late", URL: lateURL, Transport: config.MCPTransportHTTP, Enabled: true}
	store.gate, store.entered = make(chan struct{}), make(chan struct{}, 1)
	reloaded := make(chan error, 1)
	go func() { reloaded <- service.Reload(context.Background()) }()
	<-store.entered

	closed := make(chan struct{})
	go func() {
		service.Close()
		close(closed)
	}()
	select {
	case <-closed:
		t.Fatal("Close returned while a reload was still applying")
	case <-time.After(100 * time.Millisecond):
	}

	close(store.gate)
	<-closed
	require.NoError(t, <-reloaded)

	late, ok := service.manager.get("late")
	require.True(t, ok, "the in-flight reload applied before Close")
	late.stateMu.Lock()
	lateClosed := late.closed
	late.stateMu.Unlock()
	assert.True(t, lateClosed, "an upstream added by the in-flight reload is closed by Close")

	store.gate = nil
	require.ErrorIs(t, service.Reload(context.Background()), errServiceClosed, "reloads after Close do nothing")
}

func TestCloseCancelsAdminReloadBlockedOnStore(t *testing.T) {
	store := &gatedStore{memoryStore: &memoryStore{rows: map[string]ManagedServer{}}, honorCtx: true}
	service, err := NewService(context.Background(), Options{Store: store})
	require.NoError(t, err)

	// An admin save's reload, with a request context that never ends, waits
	// on a slow store read.
	store.gate, store.entered = make(chan struct{}), make(chan struct{}, 1)
	reloaded := make(chan error, 1)
	go func() { reloaded <- service.Reload(context.Background()) }()
	<-store.entered

	closed := make(chan struct{})
	go func() {
		service.Close()
		close(closed)
	}()
	select {
	case <-closed:
	case <-time.After(5 * time.Second):
		t.Fatal("Close waited on an admin reload blocked on the store")
	}
	require.Error(t, <-reloaded, "the blocked reload is cancelled")
}

func TestManagerCloseCancelsInFlightConnect(t *testing.T) {
	entered := make(chan struct{}, 1)
	cancelled := make(chan struct{})
	release := make(chan struct{})
	// An upstream that never answers its first request, so the connect hangs.
	hanging := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// The server notices a client hang-up only once the body is read.
		_, _ = io.Copy(io.Discard, r.Body)
		select {
		case entered <- struct{}{}:
		default:
		}
		select {
		case <-r.Context().Done():
			select {
			case <-cancelled:
			default:
				close(cancelled)
			}
		case <-release:
		}
	}))
	t.Cleanup(hanging.Close)
	// Registered after hanging.Close so it runs first: a still-blocked handler
	// is released before Close waits for it, even when the test fails.
	t.Cleanup(func() { close(release) })

	manager := NewManager(http.DefaultClient)
	manager.Apply([]ServerSpec{testSpec("slow", hanging.URL, nil)})
	<-entered

	manager.Close()
	select {
	case <-cancelled:
	case <-time.After(5 * time.Second):
		t.Fatal("Close did not cancel the in-flight upstream connect")
	}
}

func TestCancelledDialEndsConnectToUnresponsiveUpstream(t *testing.T) {
	release := make(chan struct{})
	entered := make(chan struct{}, 1)
	hanging := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
		select {
		case entered <- struct{}{}:
		default:
		}
		select {
		case <-r.Context().Done():
		case <-release:
		}
	}))
	t.Cleanup(hanging.Close)
	// Registered after hanging.Close so it runs first: a still-blocked handler
	// is released before Close waits for it, even when the test fails.
	t.Cleanup(func() { close(release) })

	u := newUpstream(testSpec("slow", hanging.URL, nil), http.DefaultClient)
	t.Cleanup(u.close)
	ctx, cancel := context.WithCancel(context.Background())
	refreshed := make(chan error, 1)
	go func() { refreshed <- u.refresh(ctx) }()
	<-entered

	// The SDK waits for in-flight requests when a handshake fails, so this
	// returns only because the dial's HTTP requests are cancelled with it.
	cancel()
	select {
	case err := <-refreshed:
		require.Error(t, err)
	case <-time.After(5 * time.Second):
		t.Fatal("a cancelled dial kept waiting on an unresponsive upstream")
	}
}

func TestStaleVirtualSessionClosesOnNextRequest(t *testing.T) {
	service, _, gatewayURL := newVirtualAdminTestService(t, nil,
		ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
	sessionID := initializeRawSession(t, gatewayURL+"/mcp/coding", nil)
	service.bindMu.Lock()
	host := service.bindings[sessionID].server
	service.bindMu.Unlock()
	require.NotNil(t, host)

	// The session registered after a delete's scan: the virtual server is
	// gone, but its binding and SDK session are still there.
	service.virtualMu.Lock()
	delete(service.virtualSpecs, "coding")
	service.virtualMu.Unlock()

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

	service.bindMu.Lock()
	_, bound := service.bindings[sessionID]
	service.bindMu.Unlock()
	assert.False(t, bound, "the stale binding is dropped")
	assert.Eventually(t, func() bool { return len(slices.Collect(host.Sessions())) == 0 }, 5*time.Second, 10*time.Millisecond, "the stale session is closed")
}

// Another instance sharing the database saved a row this one has not loaded
// yet. Name checks read stored rows, so the clash is still refused.
func TestNameChecksSeeRowsAnotherInstanceSaved(t *testing.T) {
	t.Run("virtual server after a stored server", func(t *testing.T) {
		store := &memoryStore{rows: map[string]ManagedServer{}}
		virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{}}
		service, _ := newTestServiceWithOptions(t, Options{
			ConfigServers: map[string]ServerSpec{"alpha": testSpec("alpha", newTestUpstream(t, "alpha", addEchoTool("echo")), nil)},
			Store:         store,
			VirtualStore:  virtualStore,
		})
		store.rows["coding"] = ManagedServer{Name: "coding", URL: "https://example.com/mcp", Transport: config.MCPTransportHTTP, Enabled: true}

		_, err := service.UpsertVirtual(context.Background(), ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
		require.ErrorIs(t, err, ErrInvalidVirtualServer)
		assert.Contains(t, err.Error(), `name "coding" is used by MCP server "coding"`)
		assert.Empty(t, virtualStore.rows)
	})

	t.Run("server after a stored virtual server", func(t *testing.T) {
		store := &memoryStore{rows: map[string]ManagedServer{}}
		virtualStore := &memoryVirtualStore{rows: map[string]ManagedVirtualServer{}}
		service, err := NewService(context.Background(), Options{Store: store, VirtualStore: virtualStore})
		require.NoError(t, err)
		t.Cleanup(service.Close)
		virtualStore.put(ManagedVirtualServer{Name: "coding", Servers: []string{"alpha"}})
		require.False(t, service.IsVirtual("coding"), "this instance has not loaded it yet")

		err = service.Upsert(context.Background(), ManagedServer{Name: "coding", URL: "https://example.com/mcp", Transport: config.MCPTransportHTTP})
		require.ErrorIs(t, err, ErrVirtualNameTaken)
		assert.Empty(t, store.rows)
	})
}