require_await

var RequireAwaitRule = rule.CreateRule(rule.Rule{
	Name: "require-await",
	Run: func(ctx rule.RuleContext, options any) rule.RuleListeners {
		var currentScope *scopeInfo

		enterFunction := func(node *ast.Node) {
			currentScope = &scopeInfo{
				hasAwait:      false,
				isAsyncYield:  false,
				functionFlags: checker.FunctionFlagsNormal,
				upper:         currentScope,
			}

			body := node.Body()
			if body != nil && (!ast.IsBlock(body) || len(body.AsBlock().Statements.Nodes) > 0) {
				currentScope.functionFlags = checker.GetFunctionFlags(node)
			}
		}

		exitFunction := func(node *ast.Node) {
			if currentScope.functionFlags&checker.FunctionFlagsAsync != 0 && !currentScope.hasAwait && (currentScope.functionFlags&checker.FunctionFlagsGenerator == 0 || !currentScope.isAsyncYield) {

				ctx.ReportNode(node, buildMissingAwaitMessage())
			}

			currentScope = currentScope.upper
		}

		markAsHasAwait := func() {
			if currentScope != nil {
				currentScope.hasAwait = true
			}
		}

		return rule.RuleListeners{

			ast.KindFunctionDeclaration:                      enterFunction,
			rule.ListenerOnExit(ast.KindFunctionDeclaration): exitFunction,
			ast.KindMethodDeclaration:                        enterFunction,
			rule.ListenerOnExit(ast.KindMethodDeclaration):   exitFunction,
			ast.KindConstructor:                              enterFunction,
			rule.ListenerOnExit(ast.KindConstructor):         exitFunction,
			ast.KindGetAccessor:                              enterFunction,
			rule.ListenerOnExit(ast.KindGetAccessor):         exitFunction,
			ast.KindSetAccessor:                              enterFunction,
			rule.ListenerOnExit(ast.KindSetAccessor):         exitFunction,
			ast.KindFunctionExpression:                       enterFunction,
			rule.ListenerOnExit(ast.KindFunctionExpression):  exitFunction,
			ast.KindArrowFunction: func(node *ast.Node) {
				enterFunction(node)

				if currentScope.functionFlags&checker.FunctionFlagsAsync == 0 {
					return
				}

				body := ast.SkipParentheses(node.Body())
				if ast.IsBlock(body) || ast.IsAwaitExpression(body) {
					return
				}

				if utils.IsThenableType(ctx.TypeChecker, body, ctx.TypeChecker.GetTypeAtLocation(body)) {
					markAsHasAwait()
				}
			},
			rule.ListenerOnExit(ast.KindArrowFunction): exitFunction,

			ast.KindAwaitExpression: func(node *ast.Node) { markAsHasAwait() },
			ast.KindForOfStatement: func(node *ast.Node) {
				if node.AsForInOrOfStatement().AwaitModifier != nil {
					markAsHasAwait()
				}
			},
			ast.KindVariableDeclarationList: func(node *ast.Node) {
				if ast.IsVarAwaitUsing(node) {
					markAsHasAwait()
				}
			},

			ast.KindYieldExpression: func(node *ast.Node) {
				if currentScope == nil || currentScope.isAsyncYield {
					return
				}
				argument := node.Expression()
				if currentScope.functionFlags&checker.FunctionFlagsGenerator == 0 || argument == nil {
					return
				}

				if ast.IsLiteralExpression(argument) {

					return
				}

				if node.AsYieldExpression().AsteriskToken == nil {
					if utils.IsThenableType(ctx.TypeChecker, argument, ctx.TypeChecker.GetTypeAtLocation(argument)) {
						currentScope.isAsyncYield = true
					}
					return
				}

				t := ctx.TypeChecker.GetTypeAtLocation(argument)
				hasAsyncYield := utils.TypeRecurser(t, func(t *checker.Type) bool {
					return utils.GetWellKnownSymbolPropertyOfType(t, "asyncIterator", ctx.TypeChecker) != nil
				})
				if hasAsyncYield {
					currentScope.isAsyncYield = true
				}
			},
			ast.KindReturnStatement: func(node *ast.Node) {
				if currentScope == nil || currentScope.hasAwait || currentScope.functionFlags&checker.FunctionFlagsAsync == 0 {
					return
				}

				expr := node.Expression()
				if expr != nil && utils.IsThenableType(ctx.TypeChecker, expr, ctx.TypeChecker.GetTypeAtLocation(expr)) {
					markAsHasAwait()
				}
			},
		}
	},
})