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

import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
import org.simantics.scl.compiler.common.names.Name;
import org.simantics.scl.compiler.common.names.Names;
import org.simantics.scl.compiler.compilation.CompilationContext;
import org.simantics.scl.compiler.constants.Constant;
import org.simantics.scl.compiler.elaboration.expressions.Case;
import org.simantics.scl.compiler.elaboration.expressions.EApply;
import org.simantics.scl.compiler.elaboration.expressions.EConstant;
import org.simantics.scl.compiler.elaboration.expressions.EError;
import org.simantics.scl.compiler.elaboration.expressions.EIf;
import org.simantics.scl.compiler.elaboration.expressions.ELambda;
import org.simantics.scl.compiler.elaboration.expressions.ELiteral;
import org.simantics.scl.compiler.elaboration.expressions.EMatch;
import org.simantics.scl.compiler.elaboration.expressions.ESimpleLambda;
import org.simantics.scl.compiler.elaboration.expressions.EVariable;
import org.simantics.scl.compiler.elaboration.expressions.Expression;
import org.simantics.scl.compiler.elaboration.expressions.Variable;
import org.simantics.scl.compiler.elaboration.java.Builtins;
import org.simantics.scl.compiler.elaboration.modules.SCLValue;
import org.simantics.scl.compiler.environment.Environment;
import org.simantics.scl.compiler.errors.ErrorLog;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.internal.codegen.types.JavaReferenceValidator;
import org.simantics.scl.compiler.internal.codegen.types.JavaTypeTranslator;
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.compiler.types.util.MultiFunction;

import gnu.trove.list.array.TLongArrayList;
import gnu.trove.map.hash.THashMap;

public class SimplificationContext implements EnvironmentalContext {
    CompilationContext compilationContext;
    Environment environment;
    ErrorLog errorLog;
    
    THashMap<Name, SCLValue> constants = new THashMap<Name, SCLValue>();
    THashMap<Variable, Expression> inlinedVariables = new THashMap<Variable, Expression>();
    
    TLongArrayList locatableStack = new TLongArrayList();
    long locatable;
    JavaTypeTranslator javaTypeTranslator;
    JavaReferenceValidator<?, ?, ?, ?> validator;
    
    public SimplificationContext(CompilationContext compilationContext, JavaReferenceValidator<?, ?, ?, ?> validator) {
        this.compilationContext = compilationContext;
        this.environment = compilationContext.environment;
        this.errorLog = compilationContext.errorLog;
        this.javaTypeTranslator = compilationContext.javaTypeTranslator;
        this.validator = validator;         
    }
    
    public Environment getEnvironment() {
        return environment;
    }
    
    public ErrorLog getErrorLog() {
        return errorLog;
    }
    
    public void pushLocation(long loc) {
        locatableStack.add(locatable);
        locatable = loc;
    }
    
    public void popLocation() {
        locatable = locatableStack.removeAt(locatableStack.size()-1);
    }
    
    public SCLValue getValue(Name name) {
        if(constants.containsKey(name))
            return constants.get(name);
        SCLValue value = environment.getValue(name);
        if(value == null)
            errorLog.log(locatable, "Couldn't find " + name + ".");
        constants.put(name, value);
        return value;
    }
    
    public Expression getConstant(Name name, Type ... typeParameters) {
        SCLValue value = getValue(name);
        if(value == null)
            return new EError(locatable);
        return new EConstant(value, typeParameters);
    }
    
    public Expression apply(Expression f, Expression ... ps) {
        Expression result = f;
        Type type = f.getType();
        for(Expression p : ps) {
            result = new EApply(locatable, result, p);
            MultiFunction mfun;
            try {
                mfun = Types.matchFunction(type, 1);
            } catch (MatchException e) {
                throw new InternalCompilerError(e);
            }
            type = mfun.returnType;
            result.setType(type);
        }
        return result;
    }
    
    public Expression tuple(Expression ... cs) {
        if(cs.length == 1)
            return cs[0];
        Type[] typeParameters = new Type[cs.length];
        for(int i=0;i<cs.length;++i)
            typeParameters[i] = cs[i].getType();
        Expression result = new EConstant(locatable, Builtins.TUPLE_CONSTRUCTORS[cs.length], typeParameters);
        for(Expression c : cs)
            result = new EApply(locatable, result, c);
        result.setType(Types.tuple(Types.getTypes(cs)));
        return result;
    }

    public Expression var(Variable var) {
        EVariable result =  new EVariable(locatable, var);
        result.setType(var.getType());
        return result;
    }
    
    public Expression simpleLambda(Variable var, Expression val) {
        ESimpleLambda result = new ESimpleLambda(var, val);
        result.setType(Types.function(var.getType(), val.getType()));
        return result;
    }
    
    public Expression lambda(Expression pat, Expression val) {
        return new ELambda(locatable,  pat, val);
    }
    
    public Expression lambda(Case ... cases) {
        return new ELambda(locatable, cases);
    }
    
    /*
    public Expression constant(SCLValue value) {
        Expression result = new EConstant(loc, value);
        result.setType(value.getType());
        return result;
    }
    
    public Expression constant(SCLValue value, Type ... typeParameters) {
        Expression result = constant(value);
        for(Type typeParameter : typeParameters)
            result = new EApplyType(loc, result, typeParameter);
        result.setType(Types.instantiate(value.getType(), typeParameters));
        return result;
    }
    */
    
    public Expression if_(Expression condition, Expression then_, Expression else_) {
        return new EIf(locatable, condition, then_, else_);
    }
    
    public Expression mapList(Expression f, Expression l) {
        try {
            MultiFunction mfun = Types.matchFunction(f.getType(), 1);
            return apply(getConstant(Names.Prelude_mapList, new Type[] {mfun.parameterTypes[0], mfun.returnType}), f, l);
        } catch (MatchException e) {
            throw new InternalCompilerError(e);
        }
    }
    
    public Expression guardList(Expression cond) {
        return apply(getConstant(Names.Prelude_guardList), cond);
    }

    public Expression concatMap(Expression f, Expression l) {
        try {
            MultiFunction mfun = Types.matchFunction(f.getType(), 1);
            return apply(getConstant(Names.Prelude_concatMap, new Type[] {
                    mfun.parameterTypes[0], mfun.effect,
                    Types.matchApply(Types.LIST, mfun.returnType)}
            ), f, l);
        } catch (MatchException e) {
            throw new InternalCompilerError(e);
        }
    }
    
    public Expression emptyList(Type type) {
        return getConstant(Names.Prelude_emptyList, type);
    }
    
    public Expression singletonList(Expression e) {
        return apply(getConstant(Names.Prelude_singletonList, e.getType()), e);
    }

    public Expression match(Expression scrutinee, Expression pattern, Expression value) {
        Case case_ = new Case(pattern, value);
        return new EMatch(scrutinee, new Case[] { case_ });        
    }
    
    public Expression match(Expression scrutinee, Case ... cases) {
        return new EMatch(scrutinee, cases);
    }

    public Expression literal(Constant constant) {
        return new ELiteral(constant);
    }
    
    @SuppressWarnings({ "unchecked" })
    public JavaReferenceValidator<Object,Object,Object,Object> getJavaReferenceValidator() {
        return (JavaReferenceValidator<Object,Object,Object,Object>)validator;
    }
    
    public JavaTypeTranslator getJavaTypeTranslator() {
        return javaTypeTranslator;
    }

    /**
     * Variable added to the context will be inlined to the
     * given expression in subsequent simplifications. It is assumed
     * that the value is already simplified.
     */
    public void addInlinedVariable(Variable variable, Expression value) {
        inlinedVariables.put(variable, value);        
    }

    public Expression getInlinedValue(Variable variable) {
        return inlinedVariables.get(variable);
    }

    public EVariable blank() {
        return new EVariable(new Variable("_"));
    }

    public Expression conditionalExecution(Expression condition, Expression continuation) {
        return new EIf(condition, continuation, new EConstant(Builtins.TUPLE_CONSTRUCTORS[0]));
    }
    
    public Expression iteratedExecution(Expression list, Variable variable, Expression continuation) {
        return new EApply(
                Locations.NO_LOCATION,
                Types.PROC,
                getConstant(Names.Prelude_iterList, variable.getType(), Types.PROC, Types.tupleConstructor(0)),
                new Expression[] {
                    new ESimpleLambda(Types.PROC, variable, continuation), 
                    list
                }
                );
    }

    public Expression iteratedVectorExecution(EApply vector, Variable variable,
            Expression continuation) {
        return new EApply(
                Locations.NO_LOCATION,
                Types.PROC,
                getConstant(Names.Vector_iterVector, variable.getType(), Types.PROC, Types.tupleConstructor(0)),
                new Expression[] {
                    new ESimpleLambda(Types.PROC, variable, continuation), 
                    vector
                }
                );
    }

    public Expression[] vars(Variable[] parameters) {
        Expression[] result = new Expression[parameters.length];
        for(int i=0;i<parameters.length;++i)
            result[i] = new EVariable(parameters[i]);
        return result;
    }

    public CompilationContext getCompilationContext() {
        return compilationContext;
    }
}