package org.simantics.scl.compiler.internal.codegen.chr;

import java.util.ArrayList;
import java.util.concurrent.atomic.AtomicReference;

import org.cojen.classfile.TypeDesc;
import org.objectweb.asm.Label;
import org.objectweb.asm.Opcodes;
import org.simantics.scl.compiler.elaboration.chr.CHRRule;
import org.simantics.scl.compiler.elaboration.chr.CHRRuleset;
import org.simantics.scl.compiler.elaboration.chr.plan.CHRSearchPlan;
import org.simantics.scl.compiler.elaboration.chr.relations.CHRConstraint;
import org.simantics.scl.compiler.internal.codegen.references.BoundVar;
import org.simantics.scl.compiler.internal.codegen.utils.ClassBuilder;
import org.simantics.scl.compiler.internal.codegen.utils.LocalVariable;
import org.simantics.scl.compiler.internal.codegen.utils.MethodBuilder;
import org.simantics.scl.compiler.internal.codegen.utils.MethodBuilderBase;

import gnu.trove.impl.PrimeFinder;
import gnu.trove.map.hash.THashMap;
import gnu.trove.procedure.TObjectObjectProcedure;
import gnu.trove.set.hash.THashSet;

public class CHRPriorityFactContainerCodeGenerator implements CHRCodeGenerationConstants {
    ClassBuilder storeClassBuilder;
    String containerClassName;
    private TypeDesc containerTypeDesc;

    private ClassBuilder classBuilder;
    
    private TypeDesc storeTypeDesc;
    
    private CHRRuleset ruleset;
    private CHRRule rule;

    public CHRPriorityFactContainerCodeGenerator(ClassBuilder storeClassBuilder, CHRRuleset ruleset, CHRRule rule) {
        this.storeClassBuilder = storeClassBuilder;
        
        this.containerClassName = storeClassBuilder.getClassName() + "$" + "CHRPriorityFactContainer" + rule.priority;
        this.containerTypeDesc = TypeDesc.forClass(containerClassName);
        this.classBuilder = new ClassBuilder(storeClassBuilder.getModuleBuilder(), Opcodes.ACC_PUBLIC, containerClassName, CHRPriorityFactContainer_name);
        
        this.storeTypeDesc = storeClassBuilder.getType();
        
        this.ruleset = ruleset;
        this.rule = rule;
    }
    
    public void generate() {
        generateFields();
        generateContructor();
        
        THashMap<CHRConstraint, ArrayList<CHRSearchPlan>> planMap = new THashMap<CHRConstraint, ArrayList<CHRSearchPlan>>(); 
        for(CHRSearchPlan plan : rule.plans) {
            ArrayList<CHRSearchPlan> list = planMap.get(plan.constraint);
            if(list == null) {
                list = new ArrayList<CHRSearchPlan>(4);
                planMap.put(plan.constraint, list);
            }
            list.add(plan);
        }
        planMap.forEachEntry(new TObjectObjectProcedure<CHRConstraint, ArrayList<CHRSearchPlan>>() {
            @Override
            public boolean execute(CHRConstraint constraint, ArrayList<CHRSearchPlan> plans) {
                for(int i=0;i<plans.size();++i)
                    generateActivate(constraint, plans.get(i), i);
                return true;
            }
        });
        generateActivate(planMap);
        
        classBuilder.getModuleBuilder().addClass(classBuilder);
    }
    
    private void generateContructor() {
        MethodBuilderBase mb = classBuilder.addConstructor(Opcodes.ACC_PUBLIC, new TypeDesc[] {storeTypeDesc});
        mb.loadThis();
        mb.loadConstant(rule.priority);
        mb.invokeSuperConstructor(new TypeDesc[] {TypeDesc.INT});
        mb.loadThis();
        mb.loadLocal(mb.getParameter(0));
        mb.storeField(containerClassName, "parent", storeTypeDesc);
        mb.returnVoid();
        mb.finish();
    }
    
    private void generateFields() {
        classBuilder.addField(Opcodes.ACC_PUBLIC, "parent", storeTypeDesc);
    }

    // protected abstract void activate(CHRContext context, CHRFact fact);
    private void generateActivate(THashMap<CHRConstraint, ArrayList<CHRSearchPlan>> planMap) {
        // @Override
        // public int activate(Object context, int priority) {
        //     return -1;
        // }

        MethodBuilderBase mb = classBuilder.addMethodBase(Opcodes.ACC_PUBLIC, "activate", TypeDesc.VOID, new TypeDesc[] {CHRContext, CHRFact});
        Label finishLabel = mb.createLabel();

        AtomicReference<Label> nextLabel = new AtomicReference<Label>();
        planMap.forEachEntry(new TObjectObjectProcedure<CHRConstraint, ArrayList<CHRSearchPlan>>() {
            @Override
            public boolean execute(CHRConstraint constraint, ArrayList<CHRSearchPlan> plans) {
                int nextPriority = constraint.nextPriority;
                constraint.nextPriority = rule.priority;
                
                Label next = nextLabel.get();
                if(next != null)
                    mb.setLocation(next);
                mb.loadLocal(mb.getParameter(1));
                mb.instanceOf(constraint.factTypeDesc);
                next = mb.createLabel();
                nextLabel.set(next);
                mb.ifZeroComparisonBranch(next, "==");
                
                for(int id=0;id<plans.size();++id) {
                    mb.loadThis();
                    mb.loadLocal(mb.getParameter(0));
                    mb.loadLocal(mb.getParameter(1));
                    mb.checkCast(constraint.factTypeDesc);
                    mb.invokeVirtual(classBuilder.getClassName(), "activate_" + constraint.name + "_" + id, TypeDesc.BOOLEAN, new TypeDesc[] {CHRContext, constraint.factTypeDesc});
                    mb.ifZeroComparisonBranch(finishLabel, "==");
                }
                
                // Add to priority queue
                if(nextPriority != Integer.MAX_VALUE) {
                    mb.loadThis();
                    mb.loadField(containerClassName, "parent", storeTypeDesc);
                    mb.loadField(storeClassBuilder.getClassName(), CHRCodeGenerationConstants.priorityName(nextPriority), CHRPriorityFactContainer);
                    mb.loadLocal(mb.getParameter(0));
                    mb.loadLocal(mb.getParameter(1));
                    mb.invokeVirtual(CHRPriorityFactContainer_name, "addFact", TypeDesc.VOID, new TypeDesc[] {CHRContext, CHRFact});
                }
                
                mb.branch(finishLabel);
                return true;
            }
        });
        {
            Label next = nextLabel.get();
            if(next != null)
                mb.setLocation(next);
        }
        
        mb.setLocation(finishLabel);
        mb.returnVoid();
        mb.finish();
    }
    
    private THashSet<BoundVar> usedParameters = new THashSet<BoundVar>();

    // protected abstract void activate(CHRContext context, CHRFact fact);

    private void generateActivate(CHRConstraint constraint, CHRSearchPlan plan, int id) {
        MethodBuilder mb = classBuilder.addMethod(Opcodes.ACC_PUBLIC, "activate_" + constraint.name + "_" + id, TypeDesc.BOOLEAN, new TypeDesc[] {CHRContext, constraint.factTypeDesc});
        LocalVariable priorityVar = new LocalVariable(0, containerTypeDesc);
        mb.loadLocal(priorityVar);
        mb.loadField(containerClassName, "parent", storeTypeDesc);
        LocalVariable parent = mb.createLocalVariable("parent", storeTypeDesc);
        mb.storeLocal(parent);
        BoundVar[] implementationParameters = plan.implementation.getParameters();
        mb.setLocalVariable(ruleset.this_, parent);
        mb.setLocalVariable(implementationParameters[0], mb.getParameter(0));
        mb.setLocalVariable(implementationParameters[1], mb.getParameter(1));

        // Set closure parameters
        usedParameters.clear();
        plan.implementation.forValRefs(valRef -> {
            if(valRef.getBinding() instanceof BoundVar)
                usedParameters.add((BoundVar)valRef.getBinding());
        });
        for(int j=0;j<ruleset.parameters.length;++j) {
            BoundVar parameter = ruleset.parameters[j];
            if(!usedParameters.contains(parameter))
                continue;
            mb.loadLocal(parent);
            mb.loadField(storeClassBuilder.getClassName(), CHRCodeGenerationConstants.parameterName(j), ruleset.parameterTypeDescs[j]);
            mb.store(parameter);
        }

        // Generate code
        //System.out.println("=== activate" + i + " ==========================================================");
        //System.out.println(plan.implementation);
        plan.implementation.markGenerateOnFly();
        plan.implementation.generateCodeWithAlreadyPreparedParameters(mb);
        mb.finish();
    }
}