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

import org.simantics.scl.compiler.elaboration.contexts.ReplaceContext;
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.environment.Environment;
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.internal.elaboration.utils.ExpressionDecorator;
import org.simantics.scl.compiler.internal.interpreted.IConstant;
import org.simantics.scl.compiler.internal.interpreted.IExpression;
import org.simantics.scl.compiler.internal.interpreted.IIf;
import org.simantics.scl.compiler.top.ExpressionInterpretationContext;
import org.simantics.scl.compiler.types.Type;
import org.simantics.scl.compiler.types.Types;
import org.simantics.scl.compiler.types.exceptions.MatchException;
import org.simantics.scl.runtime.tuple.Tuple0;

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

public class EIf extends Expression {
    public Expression condition;
    public Expression then_;
    public Expression else_; // may be null
    
    public EIf(Expression condition, Expression then_, Expression else_) {
        this.condition = condition;
        this.then_ = then_;
        this.else_ = else_;
    }

    public EIf(long loc, Expression condition, Expression then_, Expression else_) {
        super(loc);
        this.condition = condition;
        this.then_ = then_;
        this.else_ = else_;
    }

	public void collectRefs(TObjectIntHashMap<Object> allRefs, TIntHashSet refs) {
        condition.collectRefs(allRefs, refs);
        then_.collectRefs(allRefs, refs);
        if(else_ != null)
            else_.collectRefs(allRefs, refs);
    }

	@Override
	public void collectVars(TObjectIntHashMap<Variable> allVars,
	        TIntHashSet vars) {
	    condition.collectVars(allVars, vars);
        then_.collectVars(allVars, vars);
        if(else_ != null)
            else_.collectVars(allVars, vars);
	}

	@Override
	protected void updateType() throws MatchException {
	    setType(then_.getType());
	}
	
	@Override
	public IVal toVal(Environment env, CodeWriter w) {
        IVal conditionVal = condition.toVal(env, w); 
        CodeWriter joinPoint = w.createBlock(getType());
        CodeWriter thenBlock = w.createBlock();
        if(else_ != null) {
            CodeWriter elseBlock = w.createBlock();        
            w.if_(conditionVal, thenBlock.getContinuation(), elseBlock.getContinuation());
                
            IVal elseVal = else_.toVal(env, elseBlock);
            elseBlock.jump(joinPoint.getContinuation(), elseVal);
        }
        else {
            w.if_(conditionVal, thenBlock.getContinuation(), joinPoint.getContinuation());
        }
        IVal thenVal = then_.toVal(env, thenBlock);
        thenBlock.jump(joinPoint.getContinuation(), thenVal);
        w.continueAs(joinPoint);
        
        return w.getParameters()[0];
    }

    @Override
    public void collectFreeVariables(THashSet<Variable> vars) {
        condition.collectFreeVariables(vars);
        then_.collectFreeVariables(vars);
        if(else_ != null)
            else_.collectFreeVariables(vars);
    }

    @Override
    public Expression simplify(SimplificationContext context) {
        condition = condition.simplify(context);
        then_ = then_.simplify(context);
        if(else_ != null)
            else_ = else_.simplify(context);
        return this;
    }

    @Override
    public Expression resolve(TranslationContext context) {
        condition = condition.resolve(context);
        then_ = then_.resolve(context);
        if(else_ != null)
            else_ = else_.resolve(context);
        return this;
    }

    @Override
    public Expression replace(ReplaceContext context) {
        return new EIf(condition.replace(context), 
                then_.replace(context), 
                else_ == null ? null : else_.replace(context));
    }
    
    @Override
    public Expression checkBasicType(TypingContext context, Type requiredType) {
        condition = condition.checkType(context, Types.BOOLEAN);
        then_ = then_.checkType(context, requiredType);
        if(else_ != null)
            else_ = else_.checkType(context, requiredType);
        else
            context.getErrorLog().log(location, "Else branch is required because the return value of the if expression is used.");
        return this;
    }
    
    @Override
    public Expression checkIgnoredType(TypingContext context) {
        condition = condition.checkType(context, Types.BOOLEAN);
        then_ = then_.checkIgnoredType(context);
        if(else_ != null)
            else_ = else_.checkIgnoredType(context);
        return this;
    }
    
    @Override
    public Expression decorate(ExpressionDecorator decorator) {
        condition = condition.decorate(decorator);
        then_ = then_.decorate(decorator);
        if(else_ != null)
            else_ = else_.decorate(decorator);        
        return decorator.decorate(this);
    }
    
    @Override
    public boolean isEffectful() {
    	return condition.isEffectful() || then_.isEffectful() || (else_ != null && else_.isEffectful());
    }

    @Override
    public void collectEffects(THashSet<Type> effects) {
        condition.collectEffects(effects);
        then_.collectEffects(effects);
        if(else_ != null)
            else_.collectEffects(effects);
    }
    
    @Override
    public void setLocationDeep(long loc) {
        if(location == Locations.NO_LOCATION) {
            location = loc;
            condition.setLocationDeep(loc);
            then_.setLocationDeep(loc);
            if(else_ != null)
                else_.setLocationDeep(loc);
        }
    }
    
    @Override
    public void accept(ExpressionVisitor visitor) {
        visitor.visit(this);
    }
    
    @Override
    public IExpression toIExpression(ExpressionInterpretationContext target) {
        return new IIf(condition.toIExpression(target), then_.toIExpression(target), 
                else_ != null ? else_.toIExpression(target) : new IConstant(Tuple0.INSTANCE));
    }

    @Override
    public void forVariables(VariableProcedure procedure) {
        condition.forVariables(procedure);
        then_.forVariables(procedure);
        if(else_ != null)
            else_.forVariables(procedure);
    }
    @Override
    public Expression accept(ExpressionTransformer transformer) {
        return transformer.transform(this);
    }

    @Override
    public int getSyntacticFunctionArity() {
        return Math.max(then_.getSyntacticFunctionArity(), else_.getSyntacticFunctionArity());
    }
}