package dispatcher

import (
	"errors"
	"path/filepath"
	"testing"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
)

func TestSessionRemovalRevokesMCPToken(t *testing.T) {
	manager := newDispatcherTokenManager(t)
	previousManager := sessionMCPTokenManager
	sessionMCPTokenManager = func() *token.KeyManager { return manager }
	t.Cleanup(func() { sessionMCPTokenManager = previousManager })

	tokenValue, err := manager.GenerateAgentTokenForServer("session-key", "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}
	session := &interactiveSession{
		sessionKey: "session-key",
		agentID:    "agent-1",
		mcpToken:   tokenValue,
	}
	sessions := NewSessionManager()
	sessions.Set("session-key", session)

	if !sessions.RemoveIfCurrent("session-key", session) {
		t.Fatal("RemoveIfCurrent() = false, want true")
	}

	_, err = manager.ValidateAgentTokenInfo("session-key", tokenValue)
	if !errors.Is(err, token.ErrAgentTokenRevoked) {
		t.Fatalf("ValidateAgentTokenInfo() error = %v, want %v", err, token.ErrAgentTokenRevoked)
	}
}

func newDispatcherTokenManager(t *testing.T) *token.KeyManager {
	t.Helper()
	manager, err := token.NewKeyManager(filepath.Join(t.TempDir(), "config.yaml"))
	if err != nil {
		t.Fatalf("NewKeyManager() error = %v", err)
	}
	return manager
}