package mcp

import (
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"sort"
	"strings"
	"testing"

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

func TestMCPExportAllowlist(t *testing.T) {
	allowed := stringSet(
		"RegisterDiscoveryRoutes", "RegisterHandlerRoutes", "RegisterParamHandlerRoutes", "RegisterContextHandlerRoutes",
		"AccessControl",
		"MessageService", "NewMessageService", "MessageHandlerConfig", "NewMessageHandler",
		"SendScopeProvider",
		"NewAgentTeamHandler", "KBHandlerConfig", "NewKBHandler",
		"AgentTeamConnection", "MessageRuntime", "KBRouteContext", "MemorySearcher",
		"MessageToolDefs",
	)
	found := map[string]struct{}{}
	entries, err := os.ReadDir(".")
	require.NoError(t, err)
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".go") || strings.HasSuffix(entry.Name(), "_test.go") {
			continue
		}
		file, parseErr := parser.ParseFile(token.NewFileSet(), entry.Name(), nil, parser.ParseComments)
		require.NoError(t, parseErr)
		for _, declaration := range file.Decls {
			switch declaration := declaration.(type) {
			case *ast.GenDecl:
				for _, spec := range declaration.Specs {
					switch spec := spec.(type) {
					case *ast.TypeSpec:
						if spec.Name.IsExported() {
							found[spec.Name.Name] = struct{}{}
							require.NotNilf(t, declaration.Doc, "%s lacks GoDoc", spec.Name.Name)
						}
					case *ast.ValueSpec:
						for _, name := range spec.Names {
							if name.IsExported() {
								found[name.Name] = struct{}{}
							}
						}
					}
				}
			case *ast.FuncDecl:
				if declaration.Recv == nil && declaration.Name.IsExported() {
					found[declaration.Name.Name] = struct{}{}
					require.NotNilf(t, declaration.Doc, "%s lacks GoDoc", declaration.Name.Name)
				}
			}
		}
	}
	require.Equal(t, sortedSetKeys(allowed), sortedSetKeys(found))
}

func stringSet(values ...string) map[string]struct{} {
	result := make(map[string]struct{}, len(values))
	for _, value := range values {
		result[value] = struct{}{}
	}
	return result
}

func sortedSetKeys(values map[string]struct{}) []string {
	result := make([]string, 0, len(values))
	for value := range values {
		result = append(result, value)
	}
	sort.Strings(result)
	return result
}