package org.simantics.scl.compiler.internal.codegen.ssa.exits;

import java.util.ArrayList;
import java.util.Arrays;

import org.objectweb.asm.Label;
import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
import org.simantics.scl.compiler.constants.BooleanConstant;
import org.simantics.scl.compiler.constants.IntegerConstant;
import org.simantics.scl.compiler.constants.NoRepConstant;
import org.simantics.scl.compiler.internal.codegen.continuations.BranchRef;
import org.simantics.scl.compiler.internal.codegen.continuations.Cont;
import org.simantics.scl.compiler.internal.codegen.continuations.ContRef;
import org.simantics.scl.compiler.internal.codegen.references.Val;
import org.simantics.scl.compiler.internal.codegen.references.ValRef;
import org.simantics.scl.compiler.internal.codegen.ssa.SSABlock;
import org.simantics.scl.compiler.internal.codegen.ssa.SSAExit;
import org.simantics.scl.compiler.internal.codegen.ssa.SSAFunction;
import org.simantics.scl.compiler.internal.codegen.ssa.binders.ValRefBinder;
import org.simantics.scl.compiler.internal.codegen.utils.CopyContext;
import org.simantics.scl.compiler.internal.codegen.utils.MethodBuilder;
import org.simantics.scl.compiler.internal.codegen.utils.PrintingContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSASimplificationContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSAValidationContext;
import org.simantics.scl.compiler.internal.codegen.utils.ValRefVisitor;
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.TIntObjectHashMap;

public class Switch extends SSAExit implements ValRefBinder {

    ValRef scrutinee;
    BranchRef[] branches;
    
    public Switch(ValRef scrutinee, BranchRef[] branches) {
        this.scrutinee = scrutinee;
        this.branches = branches;
        scrutinee.setParent(this);
        for(BranchRef branch : branches)
            branch.cont.setParent(this);
    }
    
    public ValRef getScrutinee() {
        return scrutinee;
    }
    
    public BranchRef[] getBranches() {
        return branches;
    }
    
    private boolean isIntegerSwitch() {
        if(scrutinee.getType() != Types.INTEGER)
            return false;
        for(BranchRef branch : branches)
            if(branch.constructor != null && !(branch.constructor instanceof IntegerConstant))
                return false;
        return true;
    }
    
    private void generateIntegerSwitch(MethodBuilder mb) {
        int defaultId;
        for(defaultId=0;defaultId<branches.length-1&&branches[defaultId].constructor!=null;++defaultId);
        int[] values = new int[defaultId];
        Cont[] continuations = new Cont[defaultId+1];
        TIntObjectHashMap<Label> labelMap = new TIntObjectHashMap<Label>(defaultId); 
        for(int i=0;i<defaultId;++i) {
            int value = ((IntegerConstant)branches[i].constructor).getValue();
            values[i] = value;
            Cont cont = branches[i].cont.getBinding();
            labelMap.put(value,  mb.getLabel(cont));
            continuations[i] = cont;
        }
        Arrays.sort(values);
        Label[] labels = new Label[defaultId];
        for(int i=0;i<defaultId;++i)
            labels[i] = labelMap.get(values[i]);
        Label defaultLabel;
        {
            Cont cont = branches[defaultId].cont.getBinding();
            defaultLabel = mb.getLabel(cont);
            continuations[defaultId] = cont;
        }
        mb.push(scrutinee, Types.INTEGER);
        mb.switch_(values, labels, defaultLabel);
        for(Cont cont : continuations)
            mb.ensureExists(cont);
    }

    @Override
    public void generateCode(MethodBuilder mb) {
        if(isIntegerSwitch()) {
            generateIntegerSwitch(mb);
            return;
        }
        for(int i=0;i<branches.length;++i) {
            BranchRef branch = branches[i];
            if(branch.constructor == null)
                mb.jump(branch.cont);
            else if(i < branches.length-1) {                
                Label failure = mb.createLabel();
                branch.constructor.deconstruct(mb, scrutinee, 
                        branch.cont.getBinding(), failure);
                mb.setLocation(failure);
            }
            else
                branch.constructor.deconstruct(mb, scrutinee, 
                        branch.cont.getBinding(), null);
        }
    }

    @Override
    public void toString(PrintingContext context) {
        context.append("switch ");
        context.append(scrutinee);        
        for(BranchRef branch : branches) {
            context.append('\n');
            context.indentation();
            if(branch.constructor == null)
                context.append("otherwise");
            else
                context.append(branch.constructor.toString());
            context.append(" -> ");
            {
                Cont cont = branch.cont.getBinding();
                if(cont instanceof SSABlock) {
                    SSABlock block = (SSABlock)cont;
                    //if(cont.hasMoreThanOneOccurences()) {
                        context.append(cont);
                        context.append('\n');
                        context.addBlock(block);
                    /*}
                    else {
                        block.parametersToString(context);
                        context.append('\n');
                        block.bodyToString(context);
                    }*/
                }
                else {
                    context.append(cont);
                    context.append('\n');
                }
            }
        }
        for(SSABlock block : getSuccessors())
            context.addBlock(block);
    }

    @Override
    public void validate(SSAValidationContext context) {
        context.validate(scrutinee);
        if(scrutinee.getParent() != this)
            throw new InternalCompilerError();
        for(BranchRef branch : branches) {
            context.validate(branch.cont);
            if(branch.cont.getParent() != this)
                throw new InternalCompilerError();
        }
    }

    @Override
    public void destroy() {
        scrutinee.remove();
        for(BranchRef branch : branches)
            branch.cont.remove();
    }
    
    @Override
    public SSAExit copy(CopyContext context) {
        return new Switch(context.copy(scrutinee), 
                BranchRef.copy(context, branches));
    }

    @Override
    public void replace(TVar[] vars, Type[] replacements) {
        scrutinee.replace(vars, replacements);        
    }
    
    @Override
    public void simplify(SSASimplificationContext context) {
        if(Types.equals(scrutinee.getType(), Types.BOOLEAN)) {
            ContRef thenTarget = null;
            ContRef elseTarget = null;
            for(BranchRef branch : branches) {
                boolean used = false;
                if(branch.constructor == null) {
                    if(thenTarget == null) {
                        thenTarget = branch.cont;
                        used = true;
                    }
                    if(elseTarget == null) {
                        elseTarget = branch.cont;
                        used = true;
                    }
                }
                else if(((BooleanConstant)branch.constructor).getValue()) {
                    if(thenTarget == null) {
                        thenTarget = branch.cont;
                        used = true;
                    }
                }
                else {
                    if(elseTarget == null) {
                        elseTarget = branch.cont;
                        used = true;
                    }
                }
                if(!used)
                    branch.cont.remove();
            }
            
            // This may be possible if match compiler has
            // determined that one of the branches is not possible.
            if(elseTarget == null) 
                elseTarget = thenTarget;
            else if(thenTarget == null) 
                thenTarget = elseTarget;
            
            // Replace switch by jump or if
            SSAExit newExit;
            if(thenTarget == elseTarget) {
                scrutinee.remove();
                newExit = new Jump(thenTarget);
            }
            else {
                newExit = new If(scrutinee, 
                        thenTarget, 
                        elseTarget);
            }
            getParent().setExit(newExit);
            context.markModified("switch-to-if");
            newExit.simplify(context);
        }
        else if(branches.length == 1 && isConstructorParameterless(branches[0])) {
            scrutinee.remove();
            getParent().setExit(new Jump(branches[0].cont));
        }
    }
    
    private static boolean isConstructorParameterless(BranchRef branch) {
        return branch.constructor == null || branch.constructor instanceof NoRepConstant;
    }

    @Override
    public void collectFreeVariables(SSAFunction function,
            ArrayList<ValRef> vars) {
        scrutinee.collectFreeVariables(function, vars);
    }
    
    @Override
    public Cont addParametersInFrontOf(ContRef contRef, Val[] newParameters, Val[] oldParameters,
            Cont proxy) {
        if(proxy == null)
            proxy = contRef.getBinding().createProxy(getParent().getParent(), newParameters, oldParameters);
        ContRef proxyRef = proxy.createOccurrence();
        proxyRef.setParent(this);
        for(BranchRef branch : branches) {
            if(branch.cont == contRef) {
                branch.cont = proxyRef;
                break;
            }
        }
        return proxy;
    }

    @Override
    public SSABlock[] getSuccessors() {
        ArrayList<SSABlock> result = new ArrayList<SSABlock>(branches.length);
        for(BranchRef branch : branches) {
            Cont cont = branch.cont.getBinding();
            if(cont instanceof SSABlock)
                result.add((SSABlock)cont);
        }
        return result.toArray(new SSABlock[result.size()]);
    }

    @Override
    public void forValRefs(ValRefVisitor visitor) {
        visitor.visit(scrutinee);
    }
}