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

import java.util.ArrayList;

import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
import org.simantics.scl.compiler.compilation.CompilationContext;
import org.simantics.scl.compiler.constants.NoRepConstant;
import org.simantics.scl.compiler.elaboration.expressions.EAmbiguous;
import org.simantics.scl.compiler.elaboration.expressions.EApply;
import org.simantics.scl.compiler.elaboration.expressions.ELiteral;
import org.simantics.scl.compiler.elaboration.expressions.EPlaceholder;
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.modules.SCLValue;
import org.simantics.scl.compiler.environment.Environment;
import org.simantics.scl.compiler.errors.ErrorLog;
import org.simantics.scl.compiler.internal.elaboration.constraints.Constraint;
import org.simantics.scl.compiler.internal.elaboration.constraints.ConstraintEnvironment;
import org.simantics.scl.compiler.internal.elaboration.constraints.ConstraintSolver;
import org.simantics.scl.compiler.internal.elaboration.constraints.ExpressionAugmentation;
import org.simantics.scl.compiler.internal.elaboration.constraints.ReducedConstraints;
import org.simantics.scl.compiler.internal.elaboration.subsumption.Subsumption;
import org.simantics.scl.compiler.internal.elaboration.subsumption2.SubSolver2;
import org.simantics.scl.compiler.types.Skeletons;
import org.simantics.scl.compiler.types.TApply;
import org.simantics.scl.compiler.types.TCon;
import org.simantics.scl.compiler.types.TForAll;
import org.simantics.scl.compiler.types.TFun;
import org.simantics.scl.compiler.types.TMetaVar;
import org.simantics.scl.compiler.types.TPred;
import org.simantics.scl.compiler.types.TUnion;
import org.simantics.scl.compiler.types.TVar;
import org.simantics.scl.compiler.types.Type;
import org.simantics.scl.compiler.types.Types;
import org.simantics.scl.compiler.types.exceptions.UnificationException;
import org.simantics.scl.compiler.types.kinds.Kinds;
import org.simantics.scl.compiler.types.util.Polarity;
import org.simantics.scl.compiler.types.util.TypeUnparsingContext;

import gnu.trove.map.hash.THashMap;
import gnu.trove.set.hash.THashSet;

public class TypingContext {

    private CompilationContext compilationContext;
    
    // Subsumption
    private ArrayList<Subsumption> effectSubsumptions = new ArrayList<Subsumption>();
    private ArrayList<Subsumption> subsumptions = new ArrayList<Subsumption>();
    private ArrayList<Type> potentialSingletonEffects = new ArrayList<Type>();
    
    // Effects
    private ArrayList<Type> effectUpperBounds = new ArrayList<Type>();
    
    // Constraints
    private ArrayList<EVariable> constraintDemand = new ArrayList<EVariable>(); 
    private ArrayList<Variable[]> constraintFrames = new ArrayList<Variable[]>();

    // Patterns
    private boolean isInPattern;
    
    // Recursion
    public THashSet<SCLValue> recursiveValues;
    public ArrayList<EPlaceholder> recursiveReferences;
    
    // Overloading
    public ArrayList<EAmbiguous> overloadedExpressions = new ArrayList<EAmbiguous>(); 
    
    //TypeUnparsingContext tuc = new TypeUnparsingContext();   
    
    Environment environment;
    
    public TypingContext(CompilationContext compilationContext) {
        this.compilationContext = compilationContext;
        this.environment = compilationContext.environment;
    }

    /**
     * Records the fact that {@code a} should be a subeffect
     * of {@code b}. Only types {@code TMetaVar}, {@code TVar}, {@code TCon}, {@code TUnion}
     * are allowed.
     */
    public void subsumeEffect(long loc, Type a, Type b) {
        a = Types.canonical(a);
        if(a == Types.NO_EFFECTS)
            return;
        b = Types.canonical(b);
        if(a == b)
            return;
        if(b == Types.NO_EFFECTS) {
            try {
                Types.unify(a, Types.NO_EFFECTS);
            } catch(UnificationException e) {
                compilationContext.errorLog.log(loc, "No side-effects allowed here.");
                return;
            }
        }
        
        //System.out.println("subsumeEffect(" + a.toString(tuc) + ", " + b.toString(tuc) + ")");

        /*if(a instanceof TUnion) {
            TUnion union = (TUnion)a;
            for(Type e : union.effects)
                subsumeEffect(loc, e, b);
        }
        else*/
        effectSubsumptions.add(new Subsumption(loc, a, b));
    }
    
    public void subsume(long loc, Type a, Type b) throws UnificationException {
        a = Types.canonical(a);
        b = Types.canonical(b);
        if(a == b)
            return;
        
        //System.out.println("subsume(" + a.toString(tuc) + ", " + b.toString(tuc) + ")");
        
        if(a instanceof TMetaVar) {
            TMetaVar aVar = (TMetaVar)a;
            if(b instanceof TMetaVar) {
                Skeletons.unifySkeletons(a, b);
                subsumptions.add(new Subsumption(loc, a, b));
            }
            else {
                if(b.contains(aVar))
                    throw new UnificationException(a, b);
                aVar.setRef(createSubtypeTemplate(loc, b));
            }
            return;
        }
        else if(b instanceof TMetaVar) {
            TMetaVar bVar = (TMetaVar)b;
            if(a.contains(bVar))
                throw new UnificationException(a, b);
            bVar.setRef(createSupertypeTemplate(loc, a));
            return;
        }
        if(a instanceof TFun) {
            if(!(b instanceof TFun))
                throw new UnificationException(a, b);
            TFun aFun = (TFun)a;
            TFun bFun = (TFun)b;
            subsume(loc, bFun.domain, aFun.domain);
            subsumeEffect(loc, aFun.effect, bFun.effect);
            subsume(loc, aFun.range, bFun.range);
        }
        else if(a instanceof TApply) {
            Types.unify(a, b);
        }
        else if(a instanceof TPred) {
            Types.unify(a, b);
        }
        else // a instanceof TUnion
            throw new UnificationException(a, b);
    }
    
    /**
     * Creates a template for all subtypes of the given type
     * where all possible varying types are replaced by metavars
     * with properly set inequalities.
     */
    private Type createSubtypeTemplate(long loc, Type type) {
        type = Types.canonical(type);
        if(type instanceof TCon)
            return type;
        else if(type instanceof TApply)
            return type;
        else if(type instanceof TPred)
            return type;
        else if(type instanceof TFun) {
            TFun fun = (TFun)type;
            Type funEffect = Types.canonical(fun.effect);
            Type effect;
            if(funEffect == Types.NO_EFFECTS)
                effect = Types.NO_EFFECTS;
            else {
                effect = Types.metaVar(Kinds.EFFECT);
                effectSubsumptions.add(new Subsumption(loc, effect, funEffect));
            }            
            return Types.functionE(
                    createSupertypeTemplate(loc, fun.domain),
                    effect,
                    createSubtypeTemplate(loc, fun.range));
        }
        else if(type instanceof TMetaVar) {
            TMetaVar var = (TMetaVar)type;
            TMetaVar newVar = Types.metaVar(var.getKind());
            try {
                newVar.setSkeletonRef(var);
            } catch (UnificationException e) {
                throw new InternalCompilerError(loc, e);
            }
            subsumptions.add(new Subsumption(loc, newVar, var));
            return newVar;
        }
        else if(type instanceof TVar)
            return type;
        else if(type instanceof TUnion) {
            TUnion union = (TUnion)type;
            if(union.isMinimal())
                return type;
            TMetaVar newVar = Types.metaVar(Kinds.EFFECT);
            subsumptions.add(new Subsumption(loc, newVar, type));
            return newVar;
        }
        else if(type instanceof TForAll) {
            TForAll forAll = (TForAll)type;
            Type t = createSubtypeTemplate(loc, forAll.type);
            if(t == forAll.type)
                return type;
            else
                return Types.forAll(forAll.var, t);
        }
        else
            throw new InternalCompilerError("Unsupported type " + type + ".");
    }
    
    /**
     * Creates a template for all supertypes of the given type
     * where all possible varying types are replaced by metavars
     * with properly set inequalities.
     */
    private Type createSupertypeTemplate(long loc, Type type) {
        type = Types.canonical(type);
        if(type instanceof TCon)
            return type;
        else if(type instanceof TApply)
            return type;
        else if(type instanceof TPred)
            return type;
        else if(type instanceof TFun) {
            TFun fun = (TFun)type;
            TMetaVar effect = Types.metaVar(Kinds.EFFECT);
            effectSubsumptions.add(new Subsumption(loc, fun.effect, effect));
            return Types.functionE(
                    createSubtypeTemplate(loc, fun.domain),
                    effect,
                    createSupertypeTemplate(loc, fun.range));
        }
        else if(type instanceof TMetaVar) {
            TMetaVar var = (TMetaVar)type;
            TMetaVar newVar = Types.metaVar(var.getKind());
            try {
                newVar.setSkeletonRef(var);
            } catch (UnificationException e) {
                throw new InternalCompilerError(loc, e);
            }
            subsumptions.add(new Subsumption(loc, var, newVar));
            return newVar;
        }
        else if(type instanceof TVar)
            return type;
        else if(type instanceof TUnion) {
            TMetaVar newVar = Types.metaVar(Kinds.EFFECT);
            subsumptions.add(new Subsumption(loc, type, newVar));
            return newVar;
        }
        else if(type instanceof TForAll) {
            TForAll forAll = (TForAll)type;
            Type t = createSupertypeTemplate(loc, forAll.type);
            if(t == forAll.type)
                return type;
            else
                return Types.forAll(forAll.var, t);
        }
        else
            throw new InternalCompilerError("Unsupported type " + type + ".");
    }
    
    /**
     * Instantiates type abstractions and constraints from the value.
     */
    public Expression instantiate(Expression expr) {
        Type type = Types.canonical(expr.getType());
        while(type instanceof TForAll) {
            TForAll forAll = (TForAll)type;
            TVar var = forAll.var;
            TMetaVar mVar = Types.metaVar(var.getKind());
            type = Types.canonical(forAll.type).replace(var, mVar);
            expr = expr.applyType(mVar);
        }
        while(type instanceof TFun) {
            TFun fun = (TFun)type;
            if(fun.domain instanceof TPred) { // No need to canonicalize
                TPred pred = (TPred)fun.domain;            
                type = fun.range;
                
                long loc = expr.getLocation();
                EVariable var = new EVariable(loc, null);
                var.setType(pred);
                expr = new EApply(loc, expr, var);
                addConstraintDemand(var);
            }
            else if(fun.domain == Types.PUNIT) {                
                type = fun.range;
                
                long loc = expr.getLocation();
                declareEffect(expr.location, fun.effect);
                expr = new EApply(loc, fun.effect, expr, new ELiteral(loc, NoRepConstant.PUNIT));
            }
            else
                break;
        }
        expr.setType(type);
        return expr;
    }
    
    /**
     * Enforces {@code expr} to have type {@code b} by adding
     * type applications, lambdas and effect subsumptions.
     */
    public Expression subsume(Expression expr, Type b) {
        b = Types.canonical(b);
        /*if(b instanceof TForAll) {
            TForAll forAll = (TForAll)b;
            TVar var = forAll.var;
            TVar newVar = Types.var(var.getKind());
            b = Types.canonical(forAll.type).replace(var, newVar);
            return new ELambdaType(new TVar[] {newVar}, subsume(expr, b));
        }*/
        expr = instantiate(expr);

        try {
            subsume(expr.getLocation(), expr.getType(), b);
        } catch(UnificationException e) {
            typeError(expr.getLocation(), b, expr.getType());
        }
        return expr;
    }

    private boolean expandSubsumptions() {
        boolean nontrivialSubs = true;
        int iterationCount = 0;
        while(!subsumptions.isEmpty() && nontrivialSubs) {
            ArrayList<Subsumption> oldSubsumptions = subsumptions; 
            subsumptions = new ArrayList<Subsumption>();
            nontrivialSubs = false;
            for(Subsumption sub : oldSubsumptions) {
                Type a = sub.a = Types.canonical(sub.a);
                Type b = sub.b = Types.canonical(sub.b);
                if(b instanceof TMetaVar && a instanceof TMetaVar) {
                    subsumptions.add(sub);
                }
                else {
                    try {
                        subsume(sub.loc, a, b);
                    } catch (UnificationException e) {
                        compilationContext.errorLog.log(sub.loc, "Type " + a + " is not a subtype of " + b + ".");
                    }
                    nontrivialSubs = true;
                }
            }
            ++iterationCount;
            if(iterationCount == 5) {
                // Test that the skeletons are unifiable
                THashMap<TMetaVar,TMetaVar> metaVarMap =
                        new THashMap<TMetaVar,TMetaVar>();
                ArrayList<Subsumption> subCopies = new ArrayList<Subsumption>(subsumptions.size());
                for(Subsumption sub : subsumptions)
                    if(!isEffect(sub.a))
                        subCopies.add(new Subsumption(sub.loc,
                                sub.a.copySkeleton(metaVarMap), sub.b.copySkeleton(metaVarMap)));

                /*System.out.println("----");
                TypeUnparsingContext tuc = new TypeUnparsingContext();  
                for(Subsumption sub : subCopies)
                    System.out.println(sub.a.toString(tuc) + " == " + sub.b.toString(tuc));*/
                             
                for(Subsumption sub : subCopies)
                    try {
                        Types.unify(sub.a, sub.b);
                    } catch (UnificationException e) {
                        compilationContext.errorLog.log(sub.loc, "Unification of types failed.");
                        return false;
                    }
                
                /*System.out.println("----");
                for(Subsumption sub : subCopies)
                    System.out.println(sub.a.toString(tuc) + " === " + sub.b.toString(tuc));*/
            }
        }
        for(Subsumption sub : subsumptions)
            try {
                // FIXME This is not correct in all cases, should emit unsolved subsumptions to the type of the value
                Types.unify(sub.a, sub.b);
            } catch (UnificationException e) {
                // Should not happen. Both types should be metavars.
                throw new InternalCompilerError(e);
            }
        subsumptions.clear();
        return true;
    }
    
    private boolean isEffect(Type type) {
        if(type instanceof TMetaVar)
            return ((TMetaVar)type).getKind() == Kinds.EFFECT;
        else if(type instanceof TCon)
            return environment.getTypeDescriptor((TCon)type).getKind() == Kinds.EFFECT;
        else if(type instanceof TVar)
            return ((TVar)type).getKind() == Kinds.EFFECT;
        else if(type instanceof TUnion)
            return true;
        else
            return false;
        
    }
    
    public void solveSubsumptions(long globalLoc) {
        if(expandSubsumptions()) {
            if(!effectSubsumptions.isEmpty())
                SubSolver2.solve(compilationContext.errorLog, effectSubsumptions);
            
            if(!potentialSingletonEffects.isEmpty())
                try {
                    THashSet<TMetaVar> vars = new THashSet<TMetaVar>(4);
                    for(Type type : potentialSingletonEffects)
                        type.collectMetaVars(vars);
                    for(TMetaVar var : vars) {
                        if(var.getRef() == null) {
                            Polarity polarity = var.getPolarity();
                            if(!polarity.isNegative())

                                var.setRef(Types.NO_EFFECTS);
                        }
                    }
                } catch(UnificationException e) {
                    throw new InternalCompilerError(e);
                }
            //new SubSolver(compilationContext.errorLog, effectSubsumptions, potentialSingletonEffects, globalLoc).solve();
        }
    }
    
    public void declareEffect(long loc, Type effect) {
        subsumeEffect(loc, effect, effectUpperBounds.get(effectUpperBounds.size()-1));
    }
    
    public void pushEffectUpperBound(long loc, Type effect) {
        effectUpperBounds.add(effect);
        potentialSingletonEffects.add(effect);
    }
    
    public Type popEffectUpperBound() {
        return effectUpperBounds.remove(effectUpperBounds.size()-1);
    }

    public ErrorLog getErrorLog() {
        return compilationContext.errorLog;
    }

    public boolean isInPattern() {
        return isInPattern;
    }

    public void setInPattern(boolean isInPattern) {
        this.isInPattern = isInPattern;
    }

    public void pushConstraintFrame(Variable[] array) {
        constraintFrames.add(array);
    }

    public void popConstraintFrame() {
        constraintFrames.remove(constraintFrames.size()-1);
    }

    public void addConstraintDemand(EVariable var) {
        for(int i=constraintFrames.size()-1;i>=0;--i)
            for(Variable con : constraintFrames.get(i))
                if(Types.equals(var.getType(), con.getType())) {
                    var.setVariable(con);
                    return;
                }
        constraintDemand.add(var);
    }
    
    public Expression addConstraint(TPred pred) {
        EVariable evidence = new EVariable(null);
        evidence.setType(pred);
        addConstraintDemand(evidence);
        return evidence;
    }
    
    public Expression[] addConstraints(TPred[] preds) {
        if(preds.length == 0)
            return Expression.EMPTY_ARRAY;
        Expression[] evidences = new Expression[preds.length];
        for(int i=0;i<preds.length;++i)
            evidences[i] = addConstraint(preds[i]);
        return evidences;
    }

    public ArrayList<EVariable> getConstraintDemand() {
        return constraintDemand;
    }

    public void resetConstraintDemand() {
        constraintDemand = new ArrayList<EVariable>();
    }

    public void typeError(long loc, Type requiredType, Type givenType) {
        TypeUnparsingContext tuc = new TypeUnparsingContext();
        compilationContext.errorLog.log(loc, "Expected <" + requiredType.toString(tuc) + "> got <" + givenType.toString(tuc) + ">.");
    }

    public Expression solveConstraints(Environment environment, Expression expression) {
        ArrayList<EVariable> constraintDemand = getConstraintDemand();
        if(!constraintDemand.isEmpty()) {
            ConstraintEnvironment ce = new ConstraintEnvironment(compilationContext);
            ReducedConstraints red = ConstraintSolver.solve(
                    ce, new ArrayList<TPred>(0), constraintDemand,
                    true);                                        

            expression = ExpressionAugmentation.augmentSolved(
                    red.solvedConstraints, 
                    expression);

            for(Constraint c : red.unsolvedConstraints)
                compilationContext.errorLog.log(c.getDemandLocation(), "There is no instance for <"+c.constraint+">.");
        }
        else
            expression = expression.decomposeMatching();
        
        return expression;
    }

    public Environment getEnvironment() {
        return compilationContext.environment;
    }

    public CompilationContext getCompilationContext() {
        return compilationContext;
    }

}