package exec

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

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

func TestExportAllowlist(t *testing.T) {
	allowedTop := setOf(
		"RouteDependencies", "AgentReader", "TurnTargetReader", "LocalClient", "Catalog", "Error", "ErrorCode", "ErrorEnvelope", "Diagnostic", "Truncation",
		"RuntimeAuthorizationStore", "RuntimeAuthorizationConfig", "RuntimeCredentialLease",
		"ExtensionSafety", "ExtensionAction", "ExtensionDomain", "ExtensionSnapshot",
		"RuntimeActionSource", "RuntimeAction", "RuntimeCatalog", "RuntimeCatalogView",
		"RegisterRoutes", "ActionEndpoint", "DefaultCatalog", "ValidateWorkflowActionJSON", "NewError", "ErrorExitCode",
		"ApplyErrorBudget", "ErrorProjectorDigest", "RenderErrorText", "NewRuntimeAuthorizationStore", "NewExtensionSnapshot", "ValidateExtensionSnapshot",
		"FetchExtensionSnapshot", "NewRuntimeCatalog", "CatalogEndpoint",
		"EnvBaseURL", "EnvTokenFile", "CodeValidation", "CodeNotFound", "CodePermissionDenied", "CodeConflict",
		"CodeAgentExtensionCatalogStale", "CodeProviderAuthUnavailable", "CodeRateLimited",
		"CodeUpstreamFailure", "CodeInternal", "CodeResultBudgetExceeded",
		"RuntimeActionBuiltin", "RuntimeActionExtension", "PromptAction", "ExtensionPromptAction",
	)
	allowedMethods := setOf(
		"LocalClient.Call",
		"LocalClient.FetchRuntimeCatalog",
		"Catalog.Has", "Catalog.ValidateInput", "Catalog.RenderRootHelp", "Catalog.RenderDomainHelp",
		"Catalog.RenderActionHelp", "Catalog.RenderActionInputSchemaJSON", "Catalog.RenderActionOutputSchemaJSON",
		"Catalog.RenderErrorSchemaJSON", "Catalog.RenderActionSummary", "Catalog.PromptDomains", "Catalog.PromptCommonActions",
		"Error.Error", "Error.Unwrap",
		"agentOptionalString.IsZero",
		"RuntimeAuthorizationStore.Create", "RuntimeAuthorizationStore.Rotate", "RuntimeAuthorizationStore.RetireCredential", "RuntimeAuthorizationStore.Revoke",
		"ExtensionAction.ValidateInput",
		"ExtensionSnapshot.PromptGuidanceActions",
		"ExtensionSnapshot.ETag", "ExtensionSnapshot.Empty",
		"ExtensionSnapshot.Domains", "ExtensionSnapshot.Action", "ExtensionSnapshot.ExtensionNames",
		"ExtensionSnapshot.ActionNames", "ExtensionSnapshot.ValidateInput",
		"ExtensionSnapshot.RenderRootHelp", "ExtensionSnapshot.RenderDomainHelp",
		"ExtensionSnapshot.RenderActionHelp", "ExtensionSnapshot.RenderActionInputSchemaJSON",
		"ExtensionSnapshot.RenderActionOutputSchemaJSON", "ExtensionSnapshot.RuntimeGuidance",
		"RuntimeAction.Extension", "RuntimeCatalog.Lookup", "RuntimeCatalog.Actions",
		"RuntimeCatalog.Diagnostics", "RuntimeCatalog.ExtensionSnapshot", "RuntimeCatalog.ValidateInput", "RuntimeCatalog.View",
		"RuntimeCatalogView.Actions", "RuntimeCatalogView.Lookup", "RuntimeCatalogView.ValidateInput",
		"RuntimeCatalogView.RenderRootHelp", "RuntimeCatalogView.RenderDomainHelp",
		"RuntimeCatalogView.RenderActionHelp", "RuntimeCatalogView.RenderActionInputSchemaJSON",
		"RuntimeCatalogView.RenderActionOutputSchemaJSON", "RuntimeCatalogView.RenderErrorSchemaJSON",
	)
	allowedInterfaceMethods := map[string]map[string]struct{}{
		"AgentReader":      setOf("ListAgents", "GetAgent"),
		"TurnTargetReader": setOf("CurrentTurnDeliveryTarget"),
	}
	allowedFields := map[string]map[string]struct{}{
		"RouteDependencies": setOf("Authorizations", "Agents", "TurnTargets"),
		"LocalClient":       setOf("HTTPClient"),
		"Error":             setOf("Code", "Problem", "Hint", "Next", "Cause", "Stage", "Path", "Retryable", "ServerCode", "Diagnostics", "Truncation", "Truncated"),
		"ErrorEnvelope":     setOf("Code", "Problem", "Hint", "Next", "Metadata", "Stage", "Path", "Retryable", "ServerCode", "Diagnostics", "Truncation", "Truncated"),
		"Truncation":        setOf("DiagnosticsTruncated", "EnvelopeTruncated", "OmittedCount", "TruncatedFields", "SerializedBytes"),
		"RuntimeAuthorizationConfig": setOf(
			"ServerID", "AgentID", "SessionID", "RunID", "WorkDir", "Server",
			"ProcessGeneration", "RuntimeCatalog", "WorkflowCompaction", "ContextExpires", "TokenTTL",
		),
		"RuntimeCredentialLease": setOf("AuthorizationID", "CredentialID", "Token", "ExpiresAt"),
		"ExtensionSafety":        setOf("Class"),
		"PromptAction":           setOf("Domain", "Name", "ExampleJSON"),
		"ExtensionPromptAction":  setOf("ExtensionName", "ActionName"),
		"ExtensionAction": setOf(
			"ExtensionName", "DisplayName", "Name", "Summary",
			"Description", "Safety", "InputSchema", "OutputSchema", "Example",
		),
		"ExtensionDomain": setOf(
			"ExtensionName", "DisplayName", "Actions",
		),
		"RuntimeAction": setOf(
			"Source", "Domain", "Name", "Summary", "Description", "InputSchema",
			"OutputSchema", "Example", "Permission", "Safety",
		),
	}

	fset := token.NewFileSet()
	entries, err := os.ReadDir(".")
	require.NoError(t, err)
	var problems []string
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".go") || strings.HasSuffix(entry.Name(), "_test.go") {
			continue
		}
		path := filepath.Clean(entry.Name())
		file, parseErr := parser.ParseFile(fset, path, nil, parser.ParseComments)
		require.NoError(t, parseErr)
		for _, declaration := range file.Decls {
			switch declaration := declaration.(type) {
			case *ast.FuncDecl:
				if !declaration.Name.IsExported() {
					continue
				}
				if declaration.Recv == nil {
					checkAllowedExport(&problems, path, declaration.Name.Name, allowedTop, declaration.Doc)
					continue
				}
				receiver := receiverName(declaration.Recv.List[0].Type)
				if !ast.IsExported(receiver) {
					// An exported method on an unexported receiver cannot be
					// named by package consumers, so it is not public API.
					continue
				}
				if isRequiredJSONProtocolMethod(declaration.Name.Name) {
					continue
				}
				name := receiver + "." + declaration.Name.Name
				checkAllowedExport(&problems, path, name, allowedMethods, declaration.Doc)
			case *ast.GenDecl:
				for _, spec := range declaration.Specs {
					switch spec := spec.(type) {
					case *ast.TypeSpec:
						if !spec.Name.IsExported() {
							continue
						}
						checkAllowedExport(&problems, path, spec.Name.Name, allowedTop, firstDoc(spec.Doc, declaration.Doc))
						if fields, ok := allowedFields[spec.Name.Name]; ok {
							if value, structType := spec.Type.(*ast.StructType); structType {
								checkExportedFields(&problems, path, spec.Name.Name, value, fields)
							}
						}
						if methods, ok := allowedInterfaceMethods[spec.Name.Name]; ok {
							if value, interfaceType := spec.Type.(*ast.InterfaceType); interfaceType {
								checkInterfaceMethods(&problems, path, spec.Name.Name, value, methods)
							}
						}
					case *ast.ValueSpec:
						for _, name := range spec.Names {
							if name.IsExported() {
								if declaration.Tok != token.CONST {
									problems = append(problems, path+": unexpected exported var "+name.Name)
									continue
								}
								checkAllowedExport(
									&problems, path, name.Name, allowedTop, firstDoc(spec.Doc, declaration.Doc),
								)
							}
						}
					}
				}
			}
		}
	}
	sort.Strings(problems)
	require.Empty(t, problems, strings.Join(problems, "\n"))
}

func TestExportAllowlistRejectsUndocumentedAndUnexpectedAPI(t *testing.T) {
	allowed := setOf("Documented")
	var problems []string
	checkAllowedExport(&problems, "fixture.go", "Documented", allowed, nil)
	checkAllowedExport(
		&problems,
		"fixture.go",
		"Unexpected",
		allowed,
		&ast.CommentGroup{List: []*ast.Comment{{Text: "// Unexpected is outside the package API."}}},
	)
	require.Equal(t, []string{
		"fixture.go: missing GoDoc for Documented",
		"fixture.go: unexpected export Unexpected",
	}, problems)
}

func TestExportAllowlistRejectsUnexpectedFieldsAndInterfaceMethods(t *testing.T) {
	var problems []string
	checkExportedFields(
		&problems,
		"fixture.go",
		"Boundary",
		&ast.StructType{Fields: &ast.FieldList{List: []*ast.Field{{
			Names: []*ast.Ident{ast.NewIdent("Unexpected")},
		}}}},
		setOf("Allowed"),
	)
	checkInterfaceMethods(
		&problems,
		"fixture.go",
		"Source",
		&ast.InterfaceType{Methods: &ast.FieldList{List: []*ast.Field{{
			Names: []*ast.Ident{ast.NewIdent("Unexpected")},
		}}}},
		setOf("Allowed"),
	)
	require.Equal(t, []string{
		"fixture.go: unexpected exported field Boundary.Unexpected",
		"fixture.go: unexpected interface method Source.Unexpected",
		"fixture.go: missing GoDoc for Source.Unexpected",
	}, problems)
}

func checkAllowedExport(problems *[]string, path, name string, allowed map[string]struct{}, doc *ast.CommentGroup) {
	if _, ok := allowed[name]; !ok {
		*problems = append(*problems, path+": unexpected export "+name)
		return
	}
	docName := name
	if index := strings.LastIndexByte(name, '.'); index >= 0 {
		docName = name[index+1:]
	}
	if doc == nil || !strings.HasPrefix(strings.TrimSpace(doc.Text()), docName) {
		*problems = append(*problems, path+": missing GoDoc for "+name)
	}
}

func checkExportedFields(
	problems *[]string,
	path, owner string,
	value *ast.StructType,
	allowed map[string]struct{},
) {
	for _, field := range value.Fields.List {
		for _, name := range field.Names {
			if !name.IsExported() {
				continue
			}
			if _, ok := allowed[name.Name]; !ok {
				*problems = append(*problems, path+": unexpected exported field "+owner+"."+name.Name)
			}
		}
	}
}

func checkInterfaceMethods(
	problems *[]string,
	path, owner string,
	value *ast.InterfaceType,
	allowed map[string]struct{},
) {
	for _, field := range value.Methods.List {
		for _, name := range field.Names {
			if !name.IsExported() {
				continue
			}
			if _, ok := allowed[name.Name]; !ok {
				*problems = append(*problems, path+": unexpected interface method "+owner+"."+name.Name)
			}
			if field.Doc == nil || !strings.HasPrefix(strings.TrimSpace(field.Doc.Text()), name.Name) {
				*problems = append(*problems, path+": missing GoDoc for "+owner+"."+name.Name)
			}
		}
	}
}

func isRequiredJSONProtocolMethod(name string) bool {
	switch name {
	case "MarshalJSON", "UnmarshalJSON", "JSONSchema":
		return true
	default:
		return false
	}
}

func receiverName(expression ast.Expr) string {
	switch value := expression.(type) {
	case *ast.Ident:
		return value.Name
	case *ast.StarExpr:
		return receiverName(value.X)
	case *ast.IndexExpr:
		return receiverName(value.X)
	case *ast.IndexListExpr:
		return receiverName(value.X)
	default:
		return ""
	}
}

func firstDoc(groups ...*ast.CommentGroup) *ast.CommentGroup {
	for _, group := range groups {
		if group != nil {
			return group
		}
	}
	return nil
}

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