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

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

import org.cojen.classfile.TypeDesc;
import org.simantics.scl.compiler.common.exceptions.InternalCompilerError;
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.continuations.ReturnCont;
import org.simantics.scl.compiler.internal.codegen.references.BoundVar;
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.binders.BoundVarBinder;
import org.simantics.scl.compiler.internal.codegen.ssa.binders.FunctionBinder;
import org.simantics.scl.compiler.internal.codegen.ssa.exits.Jump;
import org.simantics.scl.compiler.internal.codegen.ssa.statements.LetApply;
import org.simantics.scl.compiler.internal.codegen.ssa.statements.LetFunctions;
import org.simantics.scl.compiler.internal.codegen.types.JavaTypeTranslator;
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.Printable;
import org.simantics.scl.compiler.internal.codegen.utils.PrintingContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSALambdaLiftingContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSASimplificationContext;
import org.simantics.scl.compiler.internal.codegen.utils.SSAValidationContext;
import org.simantics.scl.compiler.types.TVar;
import org.simantics.scl.compiler.types.Type;
import org.simantics.scl.compiler.types.Types;

public final class SSAFunction implements Printable, BoundVarBinder {
    Val target;
    
    TVar[] typeParameters;
    Type effect;
    SSABlock firstBlock;
    SSABlock lastBlock;
    ReturnCont returnCont;    
    
    FunctionBinder parent;
    SSAFunction prev;
    SSAFunction next;    
    
    public SSAFunction(TVar[] typeParameters, Type effect, Type returnType) {
        this(typeParameters, effect, new ReturnCont(returnType));
    }
        
    public SSAFunction(TVar[] typeParameters, Type effect, ReturnCont returnCont) {
        this.typeParameters = typeParameters;
        this.returnCont = returnCont;
        this.effect = Types.canonical(effect);
        returnCont.setParent(this);
    }

    public Val getTarget() {
        return target;
    }
    
    public void setTarget(Val target) {
        this.target = target;
        if(target instanceof BoundVar)
            ((BoundVar) target).parent = this;             
    }
    
    public void setTarget(BoundVar target) {
        this.target = target;
        target.parent = this;             
    }
    
    public boolean hasEffect() {
        return effect != Types.NO_EFFECTS;
    }
    
    public void addBlock(SSABlock block) {
        block.parent = this;
        if(lastBlock == null) {
            firstBlock = lastBlock = block;
            block.next = block.prev = null;            
        }
        else {
            lastBlock.next = block;
            block.prev = lastBlock;
            block.next = null;
            lastBlock = block;
        }
    }
    
    public void addBlockInFront(SSABlock block) {
        block.parent = this;
        if(firstBlock == null) {
            firstBlock = lastBlock = block;
            block.next = block.prev = null;            
        }
        else {
            firstBlock.prev = block;
            block.next = firstBlock;
            block.prev = null;
            firstBlock = block;
        }
    }

    public ReturnCont getReturnCont() {
        return returnCont;
    }
    
    public TVar[] getTypeParameters() {
        return typeParameters;
    }
    
    public SSABlock getFirstBlock() {
        return firstBlock;
    }
    
    public void generateCode(MethodBuilder mb) {
        JavaTypeTranslator tt = mb.getJavaTypeTranslator();
        for(int i=0,j=0;i<firstBlock.parameters.length;++i)
            if(!tt.toTypeDesc(firstBlock.parameters[i].getType()).equals(TypeDesc.VOID))
                mb.setLocalVariable(firstBlock.parameters[i], mb.getParameter(j++));
        for(SSABlock block = firstBlock; block != null; block = block.next)
            block.prepare(mb);
        firstBlock.generateCode(mb);
    }

    public void toString(PrintingContext context) {
        context.indentation();
        if(typeParameters.length > 0) {            
            context.append('<');
            boolean first = true;
            for(TVar var : typeParameters) {
                if(first)
                    first = false;
                else
                    context.append(',');
                context.append(var);
            }
            context.append("> ");
        }
        if(hasEffect()) {
            context.append(effect);
            context.append(" ");
        }
        context.append("RETURNS ");
        context.append(returnCont.getType());
        context.append('\n');
        
        // Print blocks
        context.pushBlockQueue();
        context.addBlock(getFirstBlock());
        while(true) {
            SSABlock block = context.pollBlock();
            if(block == null)
                break;
            block.toString(context);
        }
        context.popBlockQueue();
    }
    
    @Override
    public String toString() {
        PrintingContext context = new PrintingContext();
        toString(context);
        return context.toString();
    }
    
    public void validate(SSAValidationContext context) {
        if(target instanceof BoundVar && ((BoundVar)target).parent != this)
            throw new InternalCompilerError();
        
        // Add valid type variables
        for(TVar var : typeParameters)
            context.validTypeVariables.add(var);
        
        // Add valid variables and continuations
        context.validContinuations.add(returnCont);        
        for(SSABlock block = firstBlock; block != null; block = block.next) {            
            context.validContinuations.add(block);  
            for(BoundVar parameter : block.parameters)
                context.validBoundVariables.add(parameter);
            for(SSAStatement stat = block.firstStatement; stat != null; stat = stat.next)
                stat.addBoundVariablesTo(context);
        }

        // Validate blocks
        for(SSABlock block = firstBlock; block != null; block = block.next)
            block.validate(context);
        context.validate(returnCont);
        
        //context.reset(); // FIXME not good when there are nested functions
    }
    
    public void simplify(SSASimplificationContext context) {        
        /*if(target instanceof BoundVar) {
            tryToMakeMonadic(context);
        }*/
        for(SSABlock block = firstBlock; block != null; block = block.next)
            block.simplify(context);
        if(firstBlock == lastBlock && firstBlock.firstStatement == firstBlock.lastStatement &&
                firstBlock.firstStatement instanceof LetFunctions) {
            simplifySingleLambda(context);            
        }
    }
    
    private void simplifySingleLambda(SSASimplificationContext context) {
        LetFunctions letF = (LetFunctions)firstBlock.firstStatement;
        SSAFunction f = letF.getFirstFunction();
        if(f.getNext() != null)
            return;
        Val fVal = f.getTarget();
        if(!firstBlock.exit.isJump(getReturnCont(), fVal))
            return;
        if(fVal.hasMoreThanOneOccurences())
            return; // Possible if function is recursive and refers to itself
        if(hasEffect())
            return; // Not probably possible (?)
                
        // Transform
        for(SSABlock block = f.firstBlock; block != null; block = block.next)
            block.parent = this;
        lastBlock.next = f.firstBlock;
        f.firstBlock.prev = lastBlock;        
        lastBlock = f.lastBlock;
        
        firstBlock.firstStatement = firstBlock.lastStatement = null;
        setReturnCont(f.getReturnCont());        
        effect = f.effect;
        BoundVar[] newParameters = BoundVar.copy(f.firstBlock.parameters);
        firstBlock.setParameters(BoundVar.concat(getParameters(), newParameters));
        firstBlock.setExit(new Jump(f.firstBlock.createOccurrence(), ValRef.createOccurrences(newParameters)));
        context.markModified("SSAFunction.simplify-simple-lambda");
    }

    public void setReturnCont(ReturnCont returnCont) {
        this.returnCont = returnCont;
        returnCont.setParent(this);
    }
    
    /*public void tryToMakeMonadic(SSASimplificationContext context) {
        if(monadic)
            return;
        if(!Types.isApply(BTypes.PROC, 1, getReturnType()))
            return;        
        for(ValRef ref = target.getOccurrence(); ref != null; ref = ref.getNext()) {
            ValRefBinder parent = ref.getParent();
            if(!(parent instanceof LetApply))
                return;
            LetApply apply = (LetApply)parent;
            if(apply.getFunction() != ref)
                return;
            if(apply.getParameters().length != getParameters().length)
                return;
            if(!apply.hasEffect())
                return;
        }        
        makeMonadic(context);
    }*/
    
    /*private void makeMonadic(SSASimplificationContext context) {
        Type oldReturnType = returnCont.getType();
        Type newReturnType;
        try {
            newReturnType = Types.matchApply(BTypes.PROC, oldReturnType);
        } catch (MatchException e) {
            throw new InternalCompilerError();
        }
        
        //
        returnCont.setType(newReturnType);
        monadic = true;
        
        //
        SSABlock block = new SSABlock(oldReturnType);
        addBlock(block);
        
        returnCont.replaceWith(block);
        
        BoundVar x = new BoundVar(newReturnType);
        LetApply apply = new LetApply(x, block.parameters[0].createOccurrence());
        apply.setMonadic(true);
        block.addStatement(apply);
        
        block.setExit(new Jump(returnCont.createOccurrence(), x.createOccurrence()));
        context.markModified("SSAFunction.make-monadic");
    }*/

    public ValRef isEqualToConstant() {
        if(firstBlock.parameters.length > 0)
            return null;
        if(firstBlock != lastBlock)
            return null;
        if(firstBlock.firstStatement != null)
            return null;
        if(!(firstBlock.exit instanceof Jump))
            return null;
        Jump exit = (Jump)firstBlock.exit;
        if(exit.getTarget().getBinding() != returnCont)
            return null;
        return exit.getParameters()[0];
    }    

    public BoundVar[] getParameters() {
        return firstBlock.parameters;                
    }
    
    public Type[] getParameterTypes() {
        return Types.getTypes(firstBlock.parameters);                
    }

    public int getArity() {
        return firstBlock.parameters.length;
    }

    public void markGenerateOnFly() {
        for(SSABlock block = firstBlock; block != null; block = block.next)
            block.markGenerateOnFly();
    }

    public SSAFunction copy(CopyContext context) {
        TVar[] newTypeParameters = context.copyParameters(typeParameters);
        SSAFunction newFunction = new SSAFunction(newTypeParameters, effect, context.copy(returnCont));
        for(SSABlock block = firstBlock;
                block != null; block = block.next)
            newFunction.addBlock(context.copy(block));
        return newFunction;
    }

    public SSAFunction copy() {
        return copy(new CopyContext());
    }

    public void replace(TVar[] vars, Type[] replacements) {
        returnCont.replace(vars, replacements);
        for(SSABlock block = firstBlock;
                block != null; block = block.next)
            block.replace(vars, replacements);
    }

    public void setTypeParameters(TVar[] typeParameters) {
        this.typeParameters = typeParameters;
    }

    public Type getType() {
        Type type = returnCont.getType();
        type = Types.functionE(
                Types.getTypes(firstBlock.parameters),
                effect,
                type);
        type = Types.forAll(typeParameters, type);
        return type;
    }

    public void mergeBlocks(SSAFunction function) {
        SSABlock block = function.firstBlock;
        while(block != null) {
            SSABlock next = block.next;
            addBlock(block);
            block = next;
        }
    }

    public Type getReturnType() {
        return returnCont.getType();
    }

    public void destroy() {
        for(SSABlock block = firstBlock;
                block != null; block = block.next)
            block.destroy();
    }
    
    public void detach() {
        if(prev == null)
            parent.setFirstFunction(next);
        else
            prev.next = next;
        if(next != null)
            next.prev = prev;
    }
    
    public void remove() {
        destroy();
        detach();
    }
    
    public void addSibling(SSAFunction function) {
        function.parent = parent;
        function.next = next;
        function.prev = this;
        
        next.prev = function;
        next = function;
    }

    public SSAFunction getNext() {
        return next;
    }

    public void setParent(FunctionBinder parent) {
        this.parent = parent;
    }

    public void setPrev(SSAFunction function) {
        this.prev = function;
    }

    public void setNext(SSAFunction function) {
        this.next = function;
    }
    
    public void collectFreeVariables(ArrayList<ValRef> vars) {
        for(SSABlock block = firstBlock;
                block != null; block = block.next)
            block.collectFreeVariables(this, vars);
    }

    @Override
    public SSAFunction getParentFunction() {
        return parent.getParentFunction();
    }
    
    public FunctionBinder getParent() {
        return parent;
    }
    
    public void lambdaLift(SSALambdaLiftingContext context) {
        for(SSABlock block = firstBlock;
                block != null; block = block.next)
            block.lambdaLift(context);
    }

    public void addParametersInFront(BoundVar[] parameters) {
        Cont proxy = null;
        for(ContRef ref = firstBlock.getOccurrence(); ref != null; ref = ref.getNext())
            proxy = ref.addParametersInFront(parameters, firstBlock.parameters, proxy);
        
        firstBlock.parameters = BoundVar.concat(parameters, firstBlock.parameters);
        for(BoundVar parameter : parameters)
            parameter.parent = firstBlock;
    }
    
    public void apply(ValRef[] parameters) {
        if(parameters.length == 0)
            return;
        if(firstBlock.hasNoOccurences()) {
            BoundVar[] vars = firstBlock.getParameters();
            for(int i=0;i<parameters.length;++i)
                vars[i].replaceBy(parameters[i]);
            firstBlock.setParameters(Arrays.copyOfRange(vars, parameters.length, vars.length));
        }
        else {
            BoundVar[] newVars = new BoundVar[getArity()-parameters.length];
            SSABlock block = new SSABlock(newVars);
            block.setExit(new Jump(firstBlock.createOccurrence(), 
                    ValRef.concat(ValRef.copy(parameters), ValRef.createOccurrences(newVars))));
            addBlockInFront(block);
        }
    }

    public void applyTypes(Type[] types) {
        if(types.length == 0)
            return;
        if(types.length == typeParameters.length) {
            replace(typeParameters, types);
            typeParameters = TVar.EMPTY_ARRAY;
        }
        else {
            replace(Arrays.copyOf(typeParameters, types.length), types);
            typeParameters = 
                    Arrays.copyOfRange(typeParameters, 
                            types.length, typeParameters.length);
        }            
    }
    
    public boolean isSimpleEnoughForInline() {
        return firstBlock == lastBlock && 
                (firstBlock.firstStatement == null || 
                (firstBlock.firstStatement == firstBlock.lastStatement
                && firstBlock.firstStatement instanceof LetApply));
    }

    public Type getEffect() {
        return effect;
    }

    public int getBlockCount() {
        int count = 0;
        for(SSABlock block = firstBlock;block != null;block = block.next)
            ++count;
        return count;
    }

}