package lang

type Effect uint8

const (
	NoEffect        Effect = 0
	CoroutineEffect Effect = 1 << 0
)

func (effect Effect) Has(required Effect) bool { return effect&required != 0 }

type effectNode struct {
	key     string
	symbol  *Symbol
	direct  Effect
	callees []*Symbol
}

func inferCoroutineEffects(semantic *SemanticProgram) {
	var nodes []*effectNode
	for index := range semantic.Syntax.Functions {
		decl := &semantic.Syntax.Functions[index]
		symbol, _ := semantic.Global.Lookup(decl.Name)
		node := &effectNode{key: decl.Name, symbol: symbol}
		if decl.ExpressionBody != nil {
			collectExpressionEffects(decl.ExpressionBody, node)
		} else {
			collectFunctionEffects(decl.Body, node)
		}
		nodes = append(nodes, node)
	}
	for classIndex := range semantic.Syntax.Classes {
		class := &semantic.Syntax.Classes[classIndex]
		classSymbol := semantic.ClassInfo[class.Name]
		for methodIndex := range class.Methods {
			method := &class.Methods[methodIndex]
			node := &effectNode{key: class.Name + "." + method.Name, symbol: classSymbol.Methods[method.Name]}
			if method.ExpressionBody != nil {
				collectExpressionEffects(method.ExpressionBody, node)
			} else {
				collectFunctionEffects(method.Body, node)
			}
			nodes = append(nodes, node)
		}
	}
	changed := true
	for changed {
		changed = false
		for _, node := range nodes {
			effects := node.direct
			for _, callee := range node.callees {
				effects |= callee.Effects
			}
			if node.symbol != nil && node.symbol.Effects != effects {
				node.symbol.Effects = effects
				changed = true
			}
		}
	}
	for _, node := range nodes {
		if node.symbol != nil {
			semantic.FunctionEffects[node.key] = node.symbol.Effects
			if function, ok := node.symbol.Type.(FunctionType); ok {
				function.Effects = node.symbol.Effects
				node.symbol.Type = function
			}
		}
	}
}

func collectFunctionEffects(statements []Stmt, node *effectNode) {
	for _, statement := range statements {
		switch value := statement.(type) {
		case VarDecl:
			collectExpressionEffects(value.Value, node)
		case MultiVarDecl:
			collectExpressionEffects(value.Value, node)
		case AssignStmt:
			collectExpressionEffects(value.Value, node)
		case AddAssignStmt:
			collectExpressionEffects(value.Value, node)
		case MultiAssignStmt:
			collectExpressionEffects(value.Value, node)
		case FieldAssignStmt:
			collectExpressionEffects(value.Target, node)
			collectExpressionEffects(value.Value, node)
		case ReturnStmt:
			if value.Value != nil {
				collectExpressionEffects(value.Value, node)
			}
		case ThrowStmt:
			collectExpressionEffects(value.Value, node)
		case DeferStmt:
			collectExpressionEffects(value.Value, node)
		case ExprStmt:
			collectExpressionEffects(value.Value, node)
		case IfStmt:
			collectExpressionEffects(value.Cond, node)
			collectFunctionEffects(value.Then, node)
			collectFunctionEffects(value.Else, node)
		case WhileStmt:
			collectExpressionEffects(value.Cond, node)
			collectFunctionEffects(value.Body, node)
		case ForEachStmt:
			collectExpressionEffects(value.Source, node)
			collectFunctionEffects(value.Body, node)
		case MatchStmt:
			collectExpressionEffects(value.Value, node)
			for _, matchCase := range value.Cases {
				collectFunctionEffects(matchCase.Body, node)
			}
		case TryCatchStmt:
			collectFunctionEffects(value.TryBody, node)
			collectFunctionEffects(value.CatchBody, node)
		}
	}
}

func collectExpressionEffects(expression Expr, node *effectNode) {
	switch value := expression.(type) {
	case CallExpr:
		if ident, ok := value.Callee.(IdentExpr); ok {
			if ident.Name == "runBlocking" {
				return
			}
			if ident.Name == "withContext" {
				if len(value.Args) > 0 {
					collectExpressionEffects(value.Args[0], node)
				}
				return
			}
			if coroutineBuiltins[ident.Name] {
				node.direct |= CoroutineEffect
			}
		}
		if semantic := exprMeta(value); semantic != nil {
			if call, ok := semantic.Node.(HIRGotlinCall); ok && call.Target != nil {
				node.callees = append(node.callees, call.Target)
			}
			if call, ok := semantic.Node.(HIRGoCall); ok && call.InjectContext {
				node.direct |= CoroutineEffect
			}
		}
		collectExpressionEffects(value.Callee, node)
		for _, argument := range value.Args {
			collectExpressionEffects(argument, node)
		}
		for _, argument := range value.NamedArgs {
			collectExpressionEffects(argument.Value, node)
		}
	case UnaryExpr:
		collectExpressionEffects(value.Value, node)
	case NonNullExpr:
		collectExpressionEffects(value.Value, node)
	case TryExpr:
		collectExpressionEffects(value.Value, node)
	case BinaryExpr:
		collectExpressionEffects(value.Left, node)
		collectExpressionEffects(value.Right, node)
	case SelectorExpr:
		collectExpressionEffects(value.Receiver, node)
	case SafeSelectorExpr:
		collectExpressionEffects(value.Receiver, node)
	case IndexExpr:
		collectExpressionEffects(value.Receiver, node)
		collectExpressionEffects(value.Index, node)
	case MatchExpr:
		collectExpressionEffects(value.Value, node)
		for _, matchCase := range value.Cases {
			collectExpressionEffects(matchCase.Value, node)
		}
	case LambdaExpr:
		collectFunctionEffects(value.Body, node)
	}
}