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

import java.util.ArrayList;

import org.simantics.scl.compiler.common.names.Names;
import org.simantics.scl.compiler.compilation.CompilationContext;
import org.simantics.scl.compiler.constants.NoRepConstant;
import org.simantics.scl.compiler.elaboration.chr.CHRQuery;
import org.simantics.scl.compiler.elaboration.chr.CHRRuleset;
import org.simantics.scl.compiler.elaboration.chr.plan.PlanContext;
import org.simantics.scl.compiler.elaboration.chr.plan.PlanOp;
import org.simantics.scl.compiler.elaboration.chr.plan.PlanRealizer;
import org.simantics.scl.compiler.elaboration.chr.planning.QueryPlanningContext;
import org.simantics.scl.compiler.elaboration.contexts.SimplificationContext;
import org.simantics.scl.compiler.elaboration.contexts.TranslationContext;
import org.simantics.scl.compiler.elaboration.contexts.TypingContext;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.internal.codegen.references.IVal;
import org.simantics.scl.compiler.internal.codegen.writer.CodeWriter;
import org.simantics.scl.compiler.types.Types;
import org.simantics.scl.compiler.types.exceptions.MatchException;
import org.simantics.scl.compiler.types.kinds.Kinds;

import gnu.trove.map.hash.TObjectIntHashMap;
import gnu.trove.set.hash.THashSet;
import gnu.trove.set.hash.TIntHashSet;

public class ECHRSelect extends Expression {
    CHRQuery query;
    Variable[] existentialVariables;
    Expression expression;
    private ArrayList<PlanOp> planOps;
    private CHRRuleset currentRuleset;
    
    public ECHRSelect(Expression expression, CHRQuery query) {
        this.expression = expression;
        this.query = query;
    }

    @Override
    public void collectVars(TObjectIntHashMap<Variable> allVars, TIntHashSet vars) {
        query.collectVars(allVars, vars);
        expression.collectVars(allVars, vars);
    }

    @Override
    protected void updateType() throws MatchException {
        setType(Types.list(expression.getType()));
    }

    @Override
    public Expression inferType(TypingContext context) {
        for(Variable variable : existentialVariables)
            variable.setType(Types.metaVar(Kinds.STAR));
        query.checkType(context);
        expression = expression.inferType(context);
        return this;
    }
    
    @Override
    public Expression simplify(SimplificationContext simplificationContext) {
        this.expression = expression.simplify(simplificationContext);
        query.simplify(simplificationContext);
        
        CompilationContext compilationContext = simplificationContext.getCompilationContext();
        QueryPlanningContext context = new QueryPlanningContext(compilationContext, existentialVariables);
        if(query.createQueryPlan(context, null, -1, null))
            planOps = context.getPlanOps();

        return this;
    }
    
    @Override
    public IVal toVal(CompilationContext context, CodeWriter w) {
        IVal list = w.apply(location, context.getValue(Names.MList_create).getValue(), NoRepConstant.UNIT);
        planOps.add(new PlanOp(location) {
            @Override
            public void generateCode(CompilationContext context, PlanContext planContext, CodeWriter w) {
                w.apply(location, context.getValue(Names.MList_add).getValue(), list, expression.toVal(context, w));
            }
        });
        PlanRealizer realizer = new PlanRealizer(context, currentRuleset, currentRuleset != null ? currentRuleset.runtimeRulesetVariable : null, null, planOps);
        realizer.nextOp(w);
        return w.apply(location, context.getValue(Names.MList_freeze).getValue(), list);
    }

    @Override
    public void collectFreeVariables(THashSet<Variable> vars) {
        query.collectFreeVariables(vars);
        expression.collectFreeVariables(vars);
        if(existentialVariables != null)
            for(Variable variable : existentialVariables)
                vars.remove(variable);
    }

    @Override
    public Expression resolve(TranslationContext context) {
        currentRuleset = context.currentRuleset;
        
        context.pushExistentialFrame();
        query.resolve(context);
        context.disallowNewExistentials();
        expression = expression.resolve(context);
        existentialVariables = context.popExistentialFrame();
        return this;
    }

    @Override
    public void setLocationDeep(long loc) {
        if(location == Locations.NO_LOCATION) {
            query.setLocationDeep(loc);
            expression.setLocationDeep(loc);
        }
    }

    @Override
    public void accept(ExpressionVisitor visitor) {
        visitor.visit(this);
    }

    @Override
    public Expression accept(ExpressionTransformer transformer) {
        return transformer.transform(this);
    }
}