package lang

import "fmt"

func validateMutability(program *Program) error {
	checker := mutabilityChecker{}
	for _, fn := range program.Functions {
		if err := checker.checkFunction(fn, nil, nil); err != nil {
			return err
		}
	}
	for _, class := range program.Classes {
		fields := make(map[string]bool, len(class.Fields))
		for _, field := range class.Fields {
			fields[field.Name] = field.Mutable
		}
		for _, method := range class.Methods {
			if err := checker.checkFunction(method, fields, nil); err != nil {
				return err
			}
		}
	}
	for _, worker := range program.Workers {
		fields := make(map[string]bool, len(worker.Fields))
		for _, field := range worker.Fields {
			fields[field.Name] = field.Mutable
		}
		for _, method := range worker.Methods {
			if err := checker.checkFunction(method, nil, fields); err != nil {
				return err
			}
		}
	}
	return nil
}

type mutabilityChecker struct {
	scopes       []map[string]bool
	classFields  map[string]bool
	workerFields map[string]bool
}

func (c *mutabilityChecker) checkFunction(fn FunctionDecl, classFields map[string]bool, workerFields map[string]bool) error {
	c.scopes = nil
	c.classFields = classFields
	c.workerFields = workerFields
	c.pushScope()
	defer c.popScope()
	for _, param := range fn.Params {
		c.define(param.Name, false)
	}
	return c.checkStmts(fn.Body)
}

func (c *mutabilityChecker) checkStmts(stmts []Stmt) error {
	for _, stmt := range stmts {
		switch s := stmt.(type) {
		case VarDecl:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
			c.define(s.Name, s.Mutable)
		case MultiVarDecl:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
			for _, name := range s.Names {
				c.define(name, s.Mutable)
			}
		case AssignStmt:
			if err := c.requireMutable(s.Name, s.Pos); err != nil {
				return err
			}
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case AddAssignStmt:
			if err := c.requireMutable(s.Name, s.Pos); err != nil {
				return err
			}
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case MultiAssignStmt:
			for i, name := range s.Names {
				pos := 0
				if i < len(s.Positions) {
					pos = s.Positions[i]
				}
				if err := c.requireMutable(name, pos); err != nil {
					return err
				}
			}
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case ReturnStmt:
			if s.Value != nil {
				if err := c.checkExpr(s.Value); err != nil {
					return err
				}
			}
		case ThrowStmt:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case GoStmt:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case DeferStmt:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case ExprStmt:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
		case IfStmt:
			if err := c.checkExpr(s.Cond); err != nil {
				return err
			}
			if err := c.checkBlock(s.Then, nil); err != nil {
				return err
			}
			if err := c.checkBlock(s.Else, nil); err != nil {
				return err
			}
		case WhileStmt:
			if err := c.checkExpr(s.Cond); err != nil {
				return err
			}
			if err := c.checkBlock(s.Body, nil); err != nil {
				return err
			}
		case ForEachStmt:
			if err := c.checkExpr(s.Source); err != nil {
				return err
			}
			if err := c.checkBlock(s.Body, map[string]bool{s.Name: false}); err != nil {
				return err
			}
		case SelectStmt:
			for _, sc := range s.Cases {
				if err := c.checkExpr(sc.Source); err != nil {
					return err
				}
				if err := c.checkBlock(sc.Body, map[string]bool{"it": false}); err != nil {
					return err
				}
			}
		case MatchStmt:
			if err := c.checkExpr(s.Value); err != nil {
				return err
			}
			for _, matchCase := range s.Cases {
				bindings := map[string]bool{}
				for _, binding := range matchCase.Bindings {
					bindings[binding] = false
				}
				if err := c.checkBlock(matchCase.Body, bindings); err != nil {
					return err
				}
			}
		case TryCatchStmt:
			if err := c.checkBlock(s.TryBody, nil); err != nil {
				return err
			}
			if err := c.checkBlock(s.CatchBody, map[string]bool{s.CatchName: false}); err != nil {
				return err
			}
		}
	}
	return nil
}

func (c *mutabilityChecker) checkBlock(stmts []Stmt, bindings map[string]bool) error {
	c.pushScope()
	defer c.popScope()
	for name, mutable := range bindings {
		c.define(name, mutable)
	}
	return c.checkStmts(stmts)
}

func (c *mutabilityChecker) checkExpr(expr Expr) error {
	switch e := expr.(type) {
	case UnaryExpr:
		return c.checkExpr(e.Value)
	case BinaryExpr:
		if err := c.checkExpr(e.Left); err != nil {
			return err
		}
		return c.checkExpr(e.Right)
	case CallExpr:
		if err := c.checkExpr(e.Callee); err != nil {
			return err
		}
		for _, arg := range e.Args {
			if err := c.checkExpr(arg); err != nil {
				return err
			}
		}
		for _, arg := range e.NamedArgs {
			if err := c.checkExpr(arg.Value); err != nil {
				return err
			}
		}
	case SelectorExpr:
		return c.checkExpr(e.Receiver)
	case IndexExpr:
		if err := c.checkExpr(e.Receiver); err != nil {
			return err
		}
		return c.checkExpr(e.Index)
	case EnumVariantExpr:
		for _, value := range e.Values {
			if err := c.checkExpr(value); err != nil {
				return err
			}
		}
	case LambdaExpr:
		bindings := map[string]bool{}
		if e.ImplicitIt {
			bindings["it"] = false
		}
		for _, param := range e.Params {
			bindings[param.Name] = false
		}
		return c.checkBlock(e.Body, bindings)
	}
	return nil
}

func (c *mutabilityChecker) pushScope() {
	c.scopes = append(c.scopes, map[string]bool{})
}

func (c *mutabilityChecker) popScope() {
	if len(c.scopes) == 0 {
		return
	}
	c.scopes = c.scopes[:len(c.scopes)-1]
}

func (c *mutabilityChecker) define(name string, mutable bool) {
	if len(c.scopes) == 0 {
		c.pushScope()
	}
	c.scopes[len(c.scopes)-1][name] = mutable
}

func (c *mutabilityChecker) requireMutable(name string, pos int) error {
	for i := len(c.scopes) - 1; i >= 0; i-- {
		if mutable, ok := c.scopes[i][name]; ok {
			if mutable {
				return nil
			}
			return immutableAssignmentError(name, pos)
		}
	}
	if mutable, ok := c.classFields[name]; ok {
		if mutable {
			return nil
		}
		return immutableAssignmentError(name, pos)
	}
	if mutable, ok := c.workerFields[name]; ok {
		if mutable {
			return nil
		}
		return immutableAssignmentError(name, pos)
	}
	return nil
}

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