package lang

import "fmt"

func validateMutability(semantic *SemanticProgram) error {
	checker := mutabilityChecker{semantic: semantic}
	for index := range semantic.Syntax.Functions {
		if err := checker.checkFunction(&semantic.Syntax.Functions[index], nil); err != nil {
			return err
		}
	}
	for index := range semantic.Syntax.Classes {
		class := semantic.ClassInfo[semantic.Syntax.Classes[index].Name]
		for methodIndex := range semantic.Syntax.Classes[index].Methods {
			if err := checker.checkFunction(&semantic.Syntax.Classes[index].Methods[methodIndex], class); err != nil {
				return err
			}
		}
	}
	return nil
}

type mutabilityChecker struct {
	semantic *SemanticProgram
	scope    *Scope
	class    *ClassSymbol
}

func (checker *mutabilityChecker) checkFunction(function *FunctionDecl, class *ClassSymbol) error {
	checker.class = class
	checker.scope = NewScope(checker.semantic.Global)
	if class != nil {
		_ = checker.scope.Define(&Symbol{Name: "this", Kind: VariableSymbol, Type: ClassType{Class: class}})
	}
	for _, param := range function.Params {
		typeParams := append([]string{}, function.TypeParams...)
		if class != nil {
			typeParams = append(class.TypeParams, typeParams...)
		}
		typ, _ := checker.semantic.ResolveTypeRefWithParams(param.TypeRef, typeParams)
		_ = checker.scope.Define(&Symbol{Name: param.Name, Kind: VariableSymbol, Type: typ})
	}
	return checker.checkStmts(function.Body)
}

func (checker *mutabilityChecker) checkStmts(statements []Stmt) error {
	for _, statement := range statements {
		switch value := statement.(type) {
		case VarDecl:
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
			typ, _ := checker.semantic.ResolveTypeRef(value.TypeRef)
			_ = checker.scope.Define(&Symbol{Name: value.Name, Kind: VariableSymbol, Type: typ, Mutable: value.Mutable})
		case MultiVarDecl:
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
			for _, name := range value.Names {
				_ = checker.scope.Define(&Symbol{Name: name, Kind: VariableSymbol, Type: UnknownType{}, Mutable: value.Mutable})
			}
		case AssignStmt:
			if err := checker.requireMutable(value.Name, value.Pos); err != nil {
				return err
			}
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
		case AddAssignStmt:
			if err := checker.requireMutable(value.Name, value.Pos); err != nil {
				return err
			}
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
		case MultiAssignStmt:
			for index, name := range value.Names {
				position := 0
				if index < len(value.Positions) {
					position = value.Positions[index]
				}
				if err := checker.requireMutable(name, position); err != nil {
					return err
				}
			}
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
		case ReturnStmt:
			if value.Value != nil {
				if err := checker.checkExpr(value.Value); err != nil {
					return err
				}
			}
		case ThrowStmt:
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
		case DeferStmt:
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
		case ExprStmt:
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
		case IfStmt:
			if err := checker.checkExpr(value.Cond); err != nil {
				return err
			}
			if err := checker.checkBlock(value.Then, nil); err != nil {
				return err
			}
			if err := checker.checkBlock(value.Else, nil); err != nil {
				return err
			}
		case WhileStmt:
			if err := checker.checkExpr(value.Cond); err != nil {
				return err
			}
			if err := checker.checkBlock(value.Body, nil); err != nil {
				return err
			}
		case ForEachStmt:
			if err := checker.checkExpr(value.Source); err != nil {
				return err
			}
			if err := checker.checkBlock(value.Body, []string{value.Name}); err != nil {
				return err
			}
		case MatchStmt:
			if err := checker.checkExpr(value.Value); err != nil {
				return err
			}
			for _, matchCase := range value.Cases {
				if err := checker.checkBlock(matchCase.Body, matchCase.Bindings); err != nil {
					return err
				}
			}
		case TryCatchStmt:
			if err := checker.checkBlock(value.TryBody, nil); err != nil {
				return err
			}
			if err := checker.checkBlock(value.CatchBody, []string{value.CatchName}); err != nil {
				return err
			}
		}
	}
	return nil
}

func (checker *mutabilityChecker) checkBlock(statements []Stmt, bindings []string) error {
	previous := checker.scope
	checker.scope = NewScope(previous)
	defer func() { checker.scope = previous }()
	for _, name := range bindings {
		_ = checker.scope.Define(&Symbol{Name: name, Kind: VariableSymbol, Type: UnknownType{}})
	}
	return checker.checkStmts(statements)
}

func (checker *mutabilityChecker) checkExpr(expr Expr) error {
	switch value := expr.(type) {
	case UnaryExpr:
		return checker.checkExpr(value.Value)
	case BinaryExpr:
		if err := checker.checkExpr(value.Left); err != nil {
			return err
		}
		return checker.checkExpr(value.Right)
	case CallExpr:
		if err := checker.checkExpr(value.Callee); err != nil {
			return err
		}
		for _, argument := range value.Args {
			if err := checker.checkExpr(argument); err != nil {
				return err
			}
		}
		for _, argument := range value.NamedArgs {
			if err := checker.checkExpr(argument.Value); err != nil {
				return err
			}
		}
	case SelectorExpr:
		return checker.checkExpr(value.Receiver)
	case SafeSelectorExpr:
		return checker.checkExpr(value.Receiver)
	case NonNullExpr:
		return checker.checkExpr(value.Value)
	case TryExpr:
		return checker.checkExpr(value.Value)
	case IndexExpr:
		if err := checker.checkExpr(value.Receiver); err != nil {
			return err
		}
		return checker.checkExpr(value.Index)
	case EnumVariantExpr:
		for _, item := range value.Values {
			if err := checker.checkExpr(item); err != nil {
				return err
			}
		}
	case MatchExpr:
		if err := checker.checkExpr(value.Value); err != nil {
			return err
		}
		for _, matchCase := range value.Cases {
			if err := checker.checkExpr(matchCase.Value); err != nil {
				return err
			}
		}
	case LambdaExpr:
		previous := checker.scope
		checker.scope = NewScope(previous)
		defer func() { checker.scope = previous }()
		if value.ImplicitIt {
			_ = checker.scope.Define(&Symbol{Name: "it", Kind: VariableSymbol, Type: UnknownType{}})
		}
		for _, param := range value.Params {
			typ, _ := checker.semantic.ResolveTypeRef(param.TypeRef)
			_ = checker.scope.Define(&Symbol{Name: param.Name, Kind: VariableSymbol, Type: typ})
		}
		return checker.checkStmts(value.Body)
	}
	return nil
}

func (checker *mutabilityChecker) requireMutable(name string, position int) error {
	if symbol, ok := checker.scope.Lookup(name); ok && symbol.Kind == VariableSymbol {
		if symbol.Mutable {
			return nil
		}
		return immutableAssignmentError(name, position)
	}
	if checker.class != nil {
		if field, ok := checker.class.Fields[name]; ok {
			if field.Mutable {
				return nil
			}
			return immutableAssignmentError(name, position)
		}
	}
	return nil
}

func immutableAssignmentError(name string, position int) error {
	if position > 0 {
		return fmt.Errorf("cannot reassign immutable name %s at %d", name, position)
	}
	return fmt.Errorf("cannot reassign immutable name %s", name)
}