package org.simantics.scl.compiler.elaboration.expressions.visitors;

import org.simantics.scl.compiler.elaboration.chr.CHRLiteral;
import org.simantics.scl.compiler.elaboration.chr.CHRQuery;
import org.simantics.scl.compiler.elaboration.chr.CHRRule;
import org.simantics.scl.compiler.elaboration.chr.relations.SpecialCHRRelation;
import org.simantics.scl.compiler.elaboration.expressions.Assignment;
import org.simantics.scl.compiler.elaboration.expressions.Case;
import org.simantics.scl.compiler.elaboration.expressions.EAsPattern;
import org.simantics.scl.compiler.elaboration.expressions.EBind;
import org.simantics.scl.compiler.elaboration.expressions.ECHRSelect;
import org.simantics.scl.compiler.elaboration.expressions.ELet;
import org.simantics.scl.compiler.elaboration.expressions.ERuleset.DatalogRule;
import org.simantics.scl.compiler.elaboration.expressions.ESelect;
import org.simantics.scl.compiler.elaboration.expressions.ESimpleLambda;
import org.simantics.scl.compiler.elaboration.expressions.ESimpleLet;
import org.simantics.scl.compiler.elaboration.expressions.EVariable;
import org.simantics.scl.compiler.elaboration.expressions.EViewPattern;
import org.simantics.scl.compiler.elaboration.expressions.EWhen;
import org.simantics.scl.compiler.elaboration.expressions.Expression;
import org.simantics.scl.compiler.elaboration.expressions.Variable;
import org.simantics.scl.compiler.elaboration.expressions.list.ListAssignment;
import org.simantics.scl.compiler.elaboration.expressions.list.ListGenerator;
import org.simantics.scl.compiler.elaboration.query.QExists;

import gnu.trove.set.hash.THashSet;

public class CollectFreeVariablesVisitor extends StandardExpressionVisitor {
    private THashSet<Variable> freeVariables = new THashSet<Variable>();
    
    private StandardExpressionVisitor collectBoundVariablesVisitor = new StandardExpressionVisitor() {
        @Override
        public void visit(EVariable expression) {
            if(expression.variable != null)
                freeVariables.remove(expression.variable);
        }
        
        
        @Override
        public void visit(EAsPattern expression) {
            freeVariables.remove(expression.var);
            expression.pattern.accept(this);
        }
        
        @Override
        public void visit(EViewPattern expression) {
            expression.expression.accept(CollectFreeVariablesVisitor.this);
            expression.pattern.accept(this);
        }
    };
    
    @Override
    public void visit(EVariable expression) {
        if(expression.variable != null)
            freeVariables.add(expression.variable);
    }
    
    @Override
    public void visit(EBind expression) {
        if(expression.monadEvidence != null)
            visit(expression.monadEvidence);
        expression.in.accept(this);
        expression.value.accept(this);
        expression.pattern.accept(collectBoundVariablesVisitor);
    }
    
    @Override
    public void visit(Assignment assignment) {
        assignment.value.accept(this);
        assignment.pattern.accept(collectBoundVariablesVisitor);
    }
    
    @Override
    public void visit(ELet expression) {
        expression.in.accept(this);
        for(int i=expression.assignments.length-1;i>=0;--i)
            visit(expression.assignments[i]);
    }
    
    @Override
    public void visit(Case case_) {
        case_.value.accept(this);
        for(Expression pattern : case_.patterns)
            pattern.accept(collectBoundVariablesVisitor);
    }
    
    @Override
    public void visit(ESimpleLambda expression) {
        expression.value.accept(this);
        freeVariables.remove(expression.parameter);
    }
    
    @Override
    public void visit(ESimpleLet expression) {
        expression.in.accept(this);
        expression.value.accept(this);
        freeVariables.remove(expression.variable);
    }
    
    @Override
    public void visit(ECHRSelect expression) {
        expression.expression.accept(this);
        visit(expression.query);
        for(Variable variable : expression.existentialVariables)
            freeVariables.remove(variable);
    }
    
    @Override
    public void visit(ESelect expression) {
        expression.expression.accept(this);
        expression.query.accept(this);
        for(Variable variable : expression.variables)
            freeVariables.remove(variable);
    }
    
    @Override
    public void visit(EWhen expression) {
        expression.action.accept(this);
        expression.query.accept(this);
        for(Variable variable : expression.variables)
            freeVariables.remove(variable);
    }

    @Override
    public void visit(DatalogRule rule) {
        for(Expression parameter : rule.headParameters)
            parameter.accept(this);
        rule.body.accept(this);
        for(Variable variable : rule.variables)
            freeVariables.remove(variable);
    }
    
    @Override
    public void visit(ListGenerator qualifier) {
        qualifier.pattern.accept(collectBoundVariablesVisitor);
        qualifier.value.accept(this);
    }
    
    @Override
    public void visit(ListAssignment qualifier) {
        qualifier.pattern.accept(collectBoundVariablesVisitor);
        qualifier.value.accept(this);
    }
    
    @Override
    public void visit(CHRLiteral literal) {
        if(literal.relation == SpecialCHRRelation.ASSIGN) {
            literal.parameters[0].accept(collectBoundVariablesVisitor);
            literal.parameters[1].accept(this);
        }
        else {
            for(Expression parameter : literal.parameters)
                parameter.accept(this);
            if(literal.typeConstraintEvidenceParameters != null)
                for(Expression parameter : literal.typeConstraintEvidenceParameters)
                    parameter.accept(this);
        }
    }
    
    @Override
    public void visit(CHRQuery query) {
        for(int i=query.literals.length-1;i>=0;--i) {
            visit(query.literals[i]);
        }
    }
    
    @Override
    public void visit(CHRRule rule) {
        super.visit(rule);
        for(Variable variable : rule.existentialVariables)
            freeVariables.remove(variable);
    }
    
    @Override
    public void visit(QExists query) {
        query.query.accept(this);
        for(Variable variable : query.variables)
            freeVariables.remove(variable);
    }
    
    public THashSet<Variable> getFreeVariables() {
        return freeVariables;
    }
}