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

import java.util.ArrayList;

import org.cojen.classfile.TypeDesc;
import org.simantics.scl.compiler.compilation.CompilationContext;
import org.simantics.scl.compiler.constants.BooleanConstant;
import org.simantics.scl.compiler.constants.Constant;
import org.simantics.scl.compiler.constants.IntegerConstant;
import org.simantics.scl.compiler.constants.JavaConstructor;
import org.simantics.scl.compiler.constants.JavaMethod;
import org.simantics.scl.compiler.constants.NoRepConstant;
import org.simantics.scl.compiler.constants.generic.CallJava;
import org.simantics.scl.compiler.constants.generic.MethodRef.FieldRef;
import org.simantics.scl.compiler.constants.generic.MethodRef.ObjectMethodRef;
import org.simantics.scl.compiler.constants.generic.MethodRef.SetFieldRef;
import org.simantics.scl.compiler.elaboration.chr.analysis.CHRConstraintGGInfo;
import org.simantics.scl.compiler.elaboration.chr.analysis.UsageAnalysis;
import org.simantics.scl.compiler.elaboration.chr.plan.CHRSearchPlan;
import org.simantics.scl.compiler.elaboration.chr.plan.PlanRealizer;
import org.simantics.scl.compiler.elaboration.chr.relations.CHRConstraint;
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.elaboration.expressions.Variable;
import org.simantics.scl.compiler.elaboration.expressions.VariableProcedure;
import org.simantics.scl.compiler.elaboration.expressions.block.IncludeStatement;
import org.simantics.scl.compiler.environment.AmbiguousNameException;
import org.simantics.scl.compiler.errors.Locations;
import org.simantics.scl.compiler.internal.codegen.chr.CHRCodeGenerationConstants;
import org.simantics.scl.compiler.internal.codegen.chr.CHRRuntimeRulesetCodeGenerator;
import org.simantics.scl.compiler.internal.codegen.references.BoundVar;
import org.simantics.scl.compiler.internal.codegen.references.IVal;
import org.simantics.scl.compiler.internal.codegen.ssa.SSAFunction;
import org.simantics.scl.compiler.internal.codegen.types.StandardTypeConstructor;
import org.simantics.scl.compiler.internal.codegen.utils.Constants;
import org.simantics.scl.compiler.internal.codegen.writer.CodeWriter;
import org.simantics.scl.compiler.internal.parsing.Symbol;
import org.simantics.scl.compiler.types.TCon;
import org.simantics.scl.compiler.types.TVar;
import org.simantics.scl.compiler.types.Type;
import org.simantics.scl.compiler.types.Types;

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

public class CHRRuleset extends Symbol {
    
    public static final String INIT_CONSTRAINT = "__INIT__";
    
    public ArrayList<CHRConstraint> constraints = new ArrayList<CHRConstraint>();
    public ArrayList<CHRRule> rules = new ArrayList<CHRRule>();
    public ArrayList<IncludeStatement> includes = new ArrayList<IncludeStatement>();
    public THashMap<CHRConstraint, IncludeStatement> constraintSourceMap = new THashMap<CHRConstraint, IncludeStatement>();
    public THashMap<CHRConstraint, CHRConstraintGGInfo> activeConstraintGGInfo = new THashMap<CHRConstraint, CHRConstraintGGInfo>(); // contains also imported constraints
    public THashMap<IncludeStatement, ArrayList<CHRConstraint>> inverseActiveConstraintSourceMap = new THashMap<IncludeStatement, ArrayList<CHRConstraint>>();
    
    public boolean extensible;
    
    public CHRConstraint initConstraint;
    public int priorityCount;
    
    public int initialPriorityNumber = 0;
    
    public String runtimeRulesetClassName;
    public TCon runtimeRulesetType;
    public BoundVar runtimeRulesetVariable;
    public TypeDesc runtimeRulesetTypeDesc;
    
    public static final Constant READ_CURRENT_ID = new CallJava(TVar.EMPTY_ARRAY, Types.PROC, Types.INTEGER, new Type[] {Types.CHRContext},
            null, new FieldRef(CHRCodeGenerationConstants.CHRContext_name, "currentId", CHRRuntimeRulesetCodeGenerator.FACT_ID_TYPE), null);
    public static final Constant WRITE_CURRENT_ID = new CallJava(TVar.EMPTY_ARRAY, Types.PROC, Types.UNIT, new Type[] {Types.CHRContext, Types.INTEGER},
            null, new SetFieldRef(CHRCodeGenerationConstants.CHRContext_name, "currentId", CHRRuntimeRulesetCodeGenerator.FACT_ID_TYPE), null);
    public static final Constant GENERATE_ID = new CallJava(TVar.EMPTY_ARRAY, Types.PROC, Types.INTEGER, new Type[] {Types.CHRContext},
            null, new ObjectMethodRef(false, CHRCodeGenerationConstants.CHRContext_name, "generateId", CHRRuntimeRulesetCodeGenerator.FACT_ID_TYPE, Constants.EMPTY_TYPEDESC_ARRAY), null);
    
    // FIXME remove and change the parameter of Expression.toVal
    private CompilationContext cachedContext;
    
    // For code generation
    public CHRRulesetObject rulesetObject;
    public SSAFunction initializer;
    public SSAFunction deinitializer;
    
    public CHRRuleset() {
        initConstraint = new CHRConstraint(Locations.NO_LOCATION, INIT_CONSTRAINT, Type.EMPTY_ARRAY);
        constraints.add(initConstraint);
    }
    
    public void resolve(TranslationContext context) {
        boolean failedIncludes = false;
        for(IncludeStatement include : includes) {
            try {
                include.ruleset = context.resolveRuleset(include.name.text);
                if(include.ruleset == null) {
                    failedIncludes = true;
                    context.getErrorLog().log(include.name.location, "Couldn't resolve ruleset " + include.name + ".");
                    continue;
                }
                include.value = include.value.resolve(context);
                for(CHRConstraint constraint : include.ruleset.constraints) {
                    context.newCHRConstraint(constraint.name, constraint);
                    constraintSourceMap.put(constraint, include);
                }
            } catch (AmbiguousNameException e) {
                failedIncludes = true;
                context.getErrorLog().log(include.name.location, e.getMessage());
            }
        }
        if(failedIncludes)
            return;
        for(CHRConstraint constraint : constraints)
            context.newCHRConstraint(constraint.name, constraint);
        priorityCount = 0;
        for(CHRRule rule : rules) {
            rule.resolve(context);
            rule.priority = priorityCount++;
        }
        /*for(CHRConstraint constraint : constraints) {
            Variable newVariable = context.newVariable("claim" + constraint.factClassName);
        }*/
    }

    public void collectRefs(TObjectIntHashMap<Object> allRefs, TIntHashSet refs) {
        for(IncludeStatement include : includes)
            include.value.collectRefs(allRefs, refs);
        for(CHRRule rule : rules)
            rule.collectRefs(allRefs, refs);
    }

    public void checkType(TypingContext context) {
        for(IncludeStatement include : includes)
            include.value = include.value.checkType(context, include.ruleset.runtimeRulesetType);
        for(CHRRule rule : rules)
            rule.checkType(context);
    }

    public void collectVars(TObjectIntHashMap<Variable> allVars, TIntHashSet vars) {
        for(IncludeStatement include : includes)
            include.value.collectVars(allVars, vars);
        for(CHRRule rule : rules)
            rule.collectVars(allVars, vars);
    }

    public void forVariables(VariableProcedure procedure) {
        for(IncludeStatement include : includes)
            include.value.forVariables(procedure);
        for(CHRRule rule : rules)
            rule.forVariables(procedure);
    }

    public void collectFreeVariables(THashSet<Variable> vars) {
        for(IncludeStatement include : includes)
            include.value.collectFreeVariables(vars);
        for(CHRRule rule : rules)
            rule.collectFreeVariables(vars);
    }

    public void setLocationDeep(long loc) {
        if(location == Locations.NO_LOCATION) {
            this.location = loc;
            for(CHRRule rule : rules)
                rule.setLocationDeep(loc);
        }
    }
    
    public int getMinimumPriority(CHRConstraint constraint) {
        CHRConstraintGGInfo info = activeConstraintGGInfo.get(constraint);
        if(info == null)
            return Integer.MAX_VALUE;
        else
            return info.minimumPriority;
    }
    
    public int getAndUpdateNextPriority(CHRConstraint constraint, int nextPriority) {
        CHRConstraintGGInfo info = activeConstraintGGInfo.get(constraint);
        if(info == null)
            return Integer.MAX_VALUE;
        else {
            int result = info.nextPriority;
            info.nextPriority = nextPriority;
            return result;
        }
    }

    public void compile(SimplificationContext context) {
        initializeCodeGeneration(context.getCompilationContext());
        if(extensible)
            applyExtensibleDefaults();
        UsageAnalysis.analyzeUsage(this);
        for(CHRRule rule : rules) {
            rule.compile(context.getCompilationContext(), initConstraint);
            for(CHRSearchPlan plan : rule.plans) {
                CHRConstraint constraint = plan.constraint;
                if(!activeConstraintGGInfo.containsKey(constraint)) {
                    activeConstraintGGInfo.put(constraint, new CHRConstraintGGInfo(rule.priority));
                    IncludeStatement include = constraintSourceMap.get(constraint);
                    if(include != null) {
                        ArrayList<CHRConstraint> list = inverseActiveConstraintSourceMap.get(include);
                        if(list == null) {
                            list = new ArrayList<CHRConstraint>(4);
                            inverseActiveConstraintSourceMap.put(include, list);
                        }
                        list.add(constraint);
                    }
                }
            }
        }
        // remove init constraint if it is not useful
        if(getMinimumPriority(initConstraint) == Integer.MAX_VALUE) {
            constraints.remove(0);
            initConstraint = null;
        }
    }

    private void applyExtensibleDefaults() {
        for(CHRConstraint constraint : constraints) {
            // FIXME Too much indexing!!!
            int max = 1 << constraint.parameterTypes.length;
            for(int i=0;i<max;++i)
                constraint.getOrCreateIndex(cachedContext, i);
            constraint.setMayBeRemoved();
            /*
            constraint.getOrCreateIndex(cachedContext, 0);
            if(constraint.parameterTypes.length > 0)
                constraint.getOrCreateIndex(cachedContext, 1);*/
        }
    }

    public void simplify(SimplificationContext context) {
        for(IncludeStatement include : includes)
            include.value = include.value.simplify(context);
        for(CHRRule rule : rules)
            rule.simplify(context);
    }
    
    public void setRulesetType(TCon type, String className) {
        this.runtimeRulesetType = type;
        this.runtimeRulesetClassName = className;
    }
    
    public void initializeCodeGeneration(CompilationContext context) {
        cachedContext = context; // FIXME remove
        
        boolean createTypeDesc = false;
        if(runtimeRulesetType == null) {
            String suffix = context.namingPolicy.getFreshClosureClassNameSuffix();
            setRulesetType(Types.con(context.namingPolicy.getModuleName(), "CHR" + suffix), context.namingPolicy.getModuleClassName() + suffix);
            createTypeDesc = true;
        }
        runtimeRulesetTypeDesc = TypeDesc.forClass(runtimeRulesetClassName);
        runtimeRulesetVariable = new BoundVar(runtimeRulesetType);
        for(CHRConstraint constraint : constraints)
            constraint.initializeCodeGeneration(context, this);
        if(createTypeDesc && context.module != null) // for unit testing
            context.module.addTypeDescriptor(runtimeRulesetType.name, new StandardTypeConstructor(runtimeRulesetType, TVar.EMPTY_ARRAY, runtimeRulesetTypeDesc));
    }

    public static final Constant ACTIVATE = new JavaMethod(true, CHRCodeGenerationConstants.CHRContext_name, "activate", Types.PROC, Types.UNIT, Types.CHRContext, Types.INTEGER);
    private static final Constant CREATE_CHR_CONTEXT = new JavaConstructor("org/simantics/scl/runtime/chr/CHRContext", Types.PROC, Types.CHRContext);
    
    public IVal generateCode(CodeWriter w) {
        for(IncludeStatement include : includes) {
            include.storeVar = include.value.toVal(cachedContext.environment, w);
            initialPriorityNumber = Math.max(initialPriorityNumber, include.ruleset.initialPriorityNumber + include.ruleset.priorityCount);
        }
        CHRRulesetObject object = new CHRRulesetObject(runtimeRulesetVariable, this);
        w.defineObject(object);
        for(CHRRule rule : rules) {
            for(CHRSearchPlan plan : rule.plans) {
                /*System.out.println("    plan " + plan.priority);
                for(PlanOp planOp : plan.ops)
                    System.out.println("        " + planOp);*/
                CodeWriter methodWriter = object.createMethod(w.getModuleWriter(), TVar.EMPTY_ARRAY, Types.PROC, Types.BOOLEAN, new Type[] {Types.CHRContext, plan.constraint.factType});
                plan.implementation = methodWriter.getFunction();
                IVal[] implementationParameters = methodWriter.getParameters();
                plan.activeFact.setVal(implementationParameters[1]);
                PlanRealizer realizer = new PlanRealizer(cachedContext, this, runtimeRulesetVariable, implementationParameters[0], plan.ops);
                realizer.nextOp(methodWriter);
                if(methodWriter.isUnfinished())
                    methodWriter.return_(BooleanConstant.TRUE);
            }
        }
        if(!includes.isEmpty()) {
            {
                CodeWriter methodWriter = object.createMethod(w.getModuleWriter(), TVar.EMPTY_ARRAY, Types.PROC, Types.UNIT, new Type[] {Types.CHRContext});
                initializer = methodWriter.getFunction();
                for(IncludeStatement include : includes) {
                    methodWriter.apply(include.location,
                            new JavaMethod(true, runtimeRulesetClassName, "register", Types.PROC, Types.UNIT, new Type[] {runtimeRulesetType, Types.CHRContext, include.ruleset.runtimeRulesetType}),
                            object.getTarget(), methodWriter.getParameters()[0], include.storeVar);
                }
                methodWriter.return_(NoRepConstant.UNIT);
            }
            {
                CodeWriter methodWriter = object.createMethod(w.getModuleWriter(), TVar.EMPTY_ARRAY, Types.PROC, Types.UNIT, new Type[] {Types.CHRContext});
                deinitializer = methodWriter.getFunction();
                for(IncludeStatement include : includes) {
                    methodWriter.apply(include.location,
                            new JavaMethod(true, runtimeRulesetClassName, "unregister", Types.PROC, Types.UNIT, new Type[] {runtimeRulesetType, Types.CHRContext, include.ruleset.runtimeRulesetType}),
                            object.getTarget(), methodWriter.getParameters()[0], include.storeVar);
                }
                methodWriter.return_(NoRepConstant.UNIT);
            }
        }
        if(initConstraint != null) {
            IVal chrContext = w.apply(location, CREATE_CHR_CONTEXT);
            if(initializer != null) {
                w.apply(location,
                        new JavaMethod(true, runtimeRulesetClassName, "initialize", Types.PROC, Types.UNIT, new Type[] {runtimeRulesetType, Types.CHRContext}),
                        object.getTarget(), chrContext);
            }
            IVal initFact = w.apply(location, initConstraint.constructor, w.apply(location, GENERATE_ID, chrContext));
            w.apply(location, initConstraint.addProcedure, runtimeRulesetVariable, chrContext, initFact);
            w.apply(location, ACTIVATE, chrContext, new IntegerConstant(Integer.MAX_VALUE));
            if(deinitializer != null) {
                w.apply(location,
                        new JavaMethod(true, runtimeRulesetClassName, "deinitialize", Types.PROC, Types.UNIT, new Type[] {runtimeRulesetType, Types.CHRContext}),
                        object.getTarget(), chrContext);
            }
        }
        return runtimeRulesetVariable;
    }

    public void collectEffects(THashSet<Type> effects) {
        for(CHRRule rule : rules) {
            for(CHRLiteral literal : rule.head.literals)
                literal.collectQueryEffects(effects);
            for(CHRLiteral literal : rule.head.literals)
                literal.collectEnforceEffects(effects);
        }
    }

    public void addRule(CHRRule rule) {
        rules.add(rule);
        rule.parentRuleset = this;
    }
}