package org.simantics.scl.compiler.internal.elaboration.transformations;

import java.util.ArrayList;

import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
import org.simantics.scl.compiler.common.names.Names;
import org.simantics.scl.compiler.constants.Constant;
import org.simantics.scl.compiler.constants.StringConstant;
import org.simantics.scl.compiler.elaboration.contexts.TypingContext;
import org.simantics.scl.compiler.elaboration.expressions.EApply;
import org.simantics.scl.compiler.elaboration.expressions.EConstant;
import org.simantics.scl.compiler.elaboration.expressions.ELambda;
import org.simantics.scl.compiler.elaboration.expressions.ELiteral;
import org.simantics.scl.compiler.elaboration.expressions.ESimpleLet;
import org.simantics.scl.compiler.elaboration.expressions.EVariable;
import org.simantics.scl.compiler.elaboration.expressions.Expression;
import org.simantics.scl.compiler.elaboration.expressions.Expressions;
import org.simantics.scl.compiler.elaboration.expressions.Variable;
import org.simantics.scl.compiler.elaboration.expressions.VariableProcedure;
import org.simantics.scl.compiler.elaboration.expressions.block.LetStatement;
import org.simantics.scl.compiler.elaboration.expressions.block.Statement;
import org.simantics.scl.compiler.elaboration.modules.SCLValue;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.internal.codegen.references.IVal;
import org.simantics.scl.compiler.internal.types.HashCodeUtils;
import org.simantics.scl.compiler.types.TCon;
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.MultiApply;
import org.simantics.scl.compiler.types.util.MultiFunction;

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

public class UnifiableFactory {
    private final TypingContext context;    
    /**
     * The factory generates here the statements initializing the variables needed in unification.
     */
    private final ArrayList<Statement> mappingStatements;
    private THashMap<Type, Variable> defaultGenerators =
            new THashMap<Type, Variable>();
    
    public UnifiableFactory(TypingContext context, ArrayList<Statement> phase2Statements) {
        this.context = context;
        this.mappingStatements = phase2Statements;
    }
    
    /**
     * Converts an expression of type {@code T} to an expression of type {@code Unifiable T}
     * @param variableSet The unifiable variables
     * @param uniVariableMap A map from ordinary variables to unifiables
     * @param expression The expression that is converted
     * @return
     */
    private Expression toUnifiable(
            THashSet<Variable> variableSet,
            THashMap<Variable, Variable> uniVariableMap,
            Expression expression) {
        UnifiableRep rep = toUnifiableRep(variableSet, uniVariableMap, expression);
        return rep.toExpression();
    }
    
    static interface UnifiableRep {
        Expression toExpression();
    }
    
    class ConstantRep implements UnifiableRep {
        final Expression constant;
        public ConstantRep(Expression constant) {
            this.constant = constant;
        }
        @Override
        public Expression toExpression() {
            return Expressions.apply(context.getCompilationContext(), Types.NO_EFFECTS, Names.Unifiable_uId,
                    constant.getType(), constant);
        }
    }
    
    class PendingRep implements UnifiableRep {
        final THashSet<Variable> dependences;
        final THashMap<Variable, Variable> uniVariableMap;
        final Expression value;
        public PendingRep(THashSet<Variable> dependences, THashMap<Variable, Variable> uniVariableMap,
                Expression value) {
            this.dependences = dependences;
            this.uniVariableMap = uniVariableMap;
            this.value = value;
        }
        @Override
        public Expression toExpression() {
            Expression expression = value;
            for(Variable variable : dependences)
                expression = new ESimpleLet(
                        variable,
                        extract(variable.getType(), Expressions.var(uniVariableMap.get(variable))),
                        expression);
            return Expressions.apply(context.getCompilationContext(), Types.NO_EFFECTS, Names.Unifiable_uPending,
                    value.getType(), Expressions.computation(Types.PROC, expression));
        }
    }
    
    static class UniRep implements UnifiableRep {
        final Expression uni;
        public UniRep(Expression uni) {
            this.uni = uni;
        }
        @Override
        public Expression toExpression() {
            return uni;
        }
    }
    
    /**
     * Returns null, if does not contain variables from variableSet
     */
    private UnifiableRep toUnifiableRep(
            final THashSet<Variable> variableSet,
            final THashMap<Variable, Variable> uniVariableMap,
            Expression expression) {
        if(expression instanceof EVariable) {
            Variable variable = ((EVariable)expression).getVariable();
            if(!variableSet.contains(variable))
                return new ConstantRep(expression);
            
            Variable uniVariable = uniVariableMap.get(variable);
            if(uniVariable != null)
                return new UniRep(new EVariable(uniVariable));
            else
                return new UniRep(Expressions.apply(context.getCompilationContext(), Types.PROC, Names.Unifiable_uVar, variable.getType(), Expressions.punit()));
        }
        apply:
        if(expression instanceof EApply) {
            EApply apply = (EApply)expression;
            
            if(!(apply.getFunction() instanceof EConstant))
                break apply;
            EConstant function = (EConstant)apply.getFunction();
            
            IVal val = function.getValue().getValue();
            if(!(val instanceof Constant))
                break apply;
            Constant constant = (Constant)val;
            
            int constructorTag = constant.constructorTag();
            if(constructorTag < 0)
                break apply;
            
            int arity = constant.getArity();
            Expression[] parameters = apply.getParameters(); 
            if(arity != parameters.length)
                break apply;
            
            boolean hasUnifiableParameter = false;
            boolean hasPendingParameter = false;
            UnifiableRep[] uniParameters = new UnifiableRep[arity];
            for(int i=0;i<arity;++i) {
                UnifiableRep uRep = toUnifiableRep(variableSet, uniVariableMap, parameters[i]); 
                uniParameters[i] = uRep;
                if(uRep instanceof UniRep)
                    hasUnifiableParameter = true;
                else if(uRep instanceof PendingRep)
                    hasPendingParameter = true;
            }
            
            if(hasUnifiableParameter) {
                Expression[] tupleParameters = new Expression[arity];
                for(int i=0;i<arity;++i)
                    tupleParameters[i] = uniParameters[i].toExpression();
                Expression tuple = Expressions.tuple(tupleParameters);
                return new UniRep(Expressions.apply(context.getCompilationContext(), Types.NO_EFFECTS, Names.Unifiable_uCons,
                        expression.getType(), tuple.getType(),
                        getTag(function), tuple));
            }
            else if(hasPendingParameter) {
                THashSet<Variable> dependences = new THashSet<Variable>();
                for(UnifiableRep uRep : uniParameters)
                    if(uRep instanceof PendingRep)
                        dependences.addAll(((PendingRep)uRep).dependences);
                return new PendingRep(dependences, uniVariableMap, expression);
            }
            else
                return new ConstantRep(expression);
        }
        
        // Default action
        final THashSet<Variable> dependences = new THashSet<Variable>();
        expression.forVariables(new VariableProcedure() {
            
            @Override
            public void execute(long location, Variable variable) {
                if(variableSet.contains(variable))
                    dependences.add(variable);
            }
        });
        if(dependences.isEmpty())
            return new ConstantRep(expression);
        else
            return new PendingRep(dependences, uniVariableMap, expression);
    }

    private static class Constructor {
        final SCLValue function;
        final Type[] typeParameters;
        private int hashCode;
        
        public Constructor(SCLValue function, Type[] typeParameters) {
            this.function = function;
            this.typeParameters = typeParameters;
        }
        
        @Override
        public boolean equals(Object obj) {
            if(obj == this)
                return true;
            if(obj == null || obj.getClass() != Constructor.class)
                return false;
            Constructor other = (Constructor)obj;
            if(function != other.function)
                return false;
            return Types.equals(typeParameters, other.typeParameters);
        }
        
        @Override
        public int hashCode() {
            if(hashCode == 0) {
                int hash = HashCodeUtils.SEED;
                hash = HashCodeUtils.update(hash, function.hashCode());
                for(Type typeParameter : typeParameters)
                    hash = typeParameter.hashCode(hash);
                hashCode = hash;
            }
            return hashCode;
        }
    }
    
    private THashMap<Constructor, Variable> constructorTags =
            new THashMap<Constructor, Variable>();
    
    private Expression getTag(EConstant constructorExpr) {
        Constructor key = new Constructor(constructorExpr.getValue(),
                constructorExpr.getTypeParameters());
        Variable tag = constructorTags.get(key);
        if(tag == null) {
            SCLValue sclValue = constructorExpr.getValue();
            Constant constant = (Constant)sclValue.getValue();
            int arity = constant.getArity();
            int constructorTag = constant.constructorTag();
            MultiFunction mfun;
            try {
                mfun = Types.matchFunction(constructorExpr.getType(), arity);
            } catch (MatchException e) {
                throw new InternalCompilerError(e);
            }
            
            Type[] uniParameterTypes = new Type[arity];
            for(int i=0;i<arity;++i)
                uniParameterTypes[i] = Types.apply(Names.Unifiable_Unifiable, mfun.parameterTypes[i]);
            Type tupleType = Types.tuple(uniParameterTypes);
            
            // Destructor
            Expression destructor;
            if(sclValue.getName().module.equals("Builtin") && sclValue.getName().name.charAt(0)=='(') {
                // Tuple constructor is a special case, where we can just cast the value
                destructor = Expressions.constant(context.getCompilationContext(), Names.JavaBuiltin_unsafeCoerce, mfun.returnType, tupleType);
            }
            else {
                Variable[] parameters = new Variable[arity];
                for(int i=0;i<arity;++i)
                    parameters[i] = new Variable("p" + i, mfun.parameterTypes[i]);
                Expression pattern = new EApply(constructorExpr.copy(context), Expressions.vars(parameters));
                Expression[] tupleParameters = new Expression[arity];
                for(int i=0;i<arity;++i)
                    tupleParameters[i] = Expressions.apply(context.getCompilationContext(), Types.NO_EFFECTS, Names.Unifiable_uId,
                            parameters[i].getType(), Expressions.var(parameters[i]));
                Expression value = Expressions.tuple(tupleParameters);
                destructor = new ELambda(Locations.NO_LOCATION, pattern, value);
            }
            
            // Constructor
            Expression constructor;
            {
                Variable[] parameters = new Variable[arity];
                for(int i=0;i<arity;++i)
                    parameters[i] = new Variable("p" + i, uniParameterTypes[i]);
                Expression pattern = Expressions.tuple(Expressions.vars(parameters));
                Expression[] constructorParameters = new Expression[arity];
                for(int i=0;i<arity;++i)
                    constructorParameters[i] = extract(mfun.parameterTypes[i], Expressions.var(parameters[i]));
                Expression value = new EApply(constructorExpr.copy(context), constructorParameters);
                constructor = new ELambda(Locations.NO_LOCATION, pattern, value);
            }
            
            tag = new Variable("tag", Types.apply(Names.Unifiable_UTag, mfun.returnType, tupleType));
            mappingStatements.add(new LetStatement(new EVariable(tag), 
                    Expressions.apply(context.getCompilationContext(), Types.NO_EFFECTS, Names.Unifiable_uTag, tupleType, mfun.returnType,
                            Expressions.integer(constructorTag), constructor, destructor)));
            constructorTags.put(key, tag);
        }
        return new EVariable(tag);
    }
    
    private Expression extract(Type type, Expression uni) {
        return Expressions.apply(context.getCompilationContext(), Types.PROC, Names.Unifiable_extractWithDefault,
                type, getDefaultGenerator(type), uni);
    }
    
    /**
     * Returns for the given type {@code T} a generator
     * of type {@code <Proc> T} that generates the values of the type.
     */
    private Expression getDefaultGenerator(Type type) {
        Variable generator = defaultGenerators.get(type);
        if(generator == null) {
            generator = new Variable("defGen", Types.functionE(Types.PUNIT, Types.PROC, type));
            mappingStatements.add(new LetStatement(new EVariable(generator),
                    Expressions.computation(Types.PROC, createGenerationExpression(type))
                    ));
            defaultGenerators.put(type, generator);
        }
        return new EVariable(generator);
    }
    
    private Expression createGenerationExpression(Type type) {
        MultiApply apply = Types.matchApply(type);
        //System.out.println("createGenerationExpression(" + type.toString(tuc) + ")");
        if(apply.constructor instanceof TCon) {
            if(apply.constructor.equals(Types.RESOURCE))
                return Expressions.apply(context.getCompilationContext(), Types.PROC, Names.Simantics_DB_newResource, Expressions.tuple());
            
            if(apply.constructor.equals(Types.STRING))
                return new ELiteral(new StringConstant("")); // FIXME
            
            if(apply.constructor.equals(Names.Data_XML_Element))
                return Expressions.apply(context.getCompilationContext(), Types.PROC, Names.Data_XML_createElement, Expressions.string("NO-NAME"));
            
            TCon con = (TCon)apply.constructor;
            if(con.name.charAt(0) == '(') { // (), (,), (,,),...
                int arity = con.name.length()-1;
                if(arity == 1)
                    arity = 0;
                if(arity != apply.parameters.length)
                    throw new InternalCompilerError();
                Expression[] parameters = new Expression[arity];
                for(int i=0;i<arity;++i)
                    parameters[i] = new EApply(Locations.NO_LOCATION, Types.PROC,
                            getDefaultGenerator(apply.parameters[i]), Expressions.punit());
                return Expressions.tuple(parameters);
            }
        }
        return Expressions.apply(context.getCompilationContext(), Types.NO_EFFECTS, Names.Builtin_fail,
                new ELiteral(new StringConstant("Cannot generated default instance for type " + type + ".")));
    }
        
    public Expression generateDefaultValue(Type type) {
        return Expressions.apply(Types.PROC, getDefaultGenerator(type), Expressions.punit());
    }
    
    public Expression getFromUMap(Expression umap, Expression key, Type valueType) {
        return Expressions.apply(context.getCompilationContext(), Types.PROC,
                Names.Unifiable_getUMapWithDefault,
                valueType,
                key.getType(),
                getDefaultGenerator(valueType),
                umap,
                key);
    }

    public Expression putToUMapUnifiable(
            THashSet<Variable> variableSet, THashMap<Variable, Variable> uniVariableMap,
            Expression umap, Expression key, Expression value) {
        return Expressions.apply(context.getCompilationContext(), Types.PROC,
                Names.Unifiable_putUMap,
                key.getType(),
                value.getType(),
                umap,
                key,
                toUnifiable(variableSet, uniVariableMap, value));
    }

    public Expression putToUMapConstant(Variable umap, Expression key, Expression value) {
        return Expressions.apply(context.getCompilationContext(), Types.PROC,
                Names.Unifiable_putUMapC,
                key.getType(), value.getType(),
                Expressions.var(umap),
                key, value);
    }
}