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

import java.util.Arrays;
import java.util.Map;

import org.cojen.classfile.TypeDesc;
import org.objectweb.asm.Opcodes;
import org.simantics.scl.compiler.constants.FunctionValue;
import org.simantics.scl.compiler.constants.LocalFieldConstant;
import org.simantics.scl.compiler.constants.LocalVariableConstant;
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.types.JavaTypeTranslator;
import org.simantics.scl.compiler.top.SCLCompilerConfiguration;
import org.simantics.scl.compiler.types.Type;

import gnu.trove.map.hash.THashMap;

public class ModuleBuilder {
    JavaNamingPolicy namingPolicy;
    JavaTypeTranslator javaTypeTranslator;
    
    THashMap<String, byte[]> classes = new THashMap<String, byte[]>(); 

    THashMap<ClosureDesc, TypeDesc> cache = new THashMap<ClosureDesc, TypeDesc>();
    
    MethodSizeCounter methodSizeCounter;
    
    public void addClass(ClassBuilder cb) {
        byte[] bytecode = cb.finishClass();
        classes.put(cb.getClassName(), bytecode);
        //System.out.println("Added " + cb.getClassName());
    }
    
    public JavaTypeTranslator getJavaTypeTranslator() {
        return javaTypeTranslator;
    }
        
    static class ClosureDesc {
        FunctionValue functionValue;
        int knownParametersCount;
        
        public ClosureDesc(FunctionValue functionValue, int knownParametersCount) {
            this.functionValue = functionValue;
            this.knownParametersCount = knownParametersCount;
        }
        
        @Override
        public int hashCode() {
            return functionValue.hashCode() + 31 * knownParametersCount;
        }
        
        @Override
        public boolean equals(Object obj) {
            if(this == obj)
                return true;
            if(obj == null || obj.getClass() != getClass())
                return false;
            ClosureDesc other = (ClosureDesc)obj;
            return functionValue == other.functionValue && knownParametersCount == other.knownParametersCount;
        }        
    }
    
    public ModuleBuilder(JavaNamingPolicy namingPolicy, JavaTypeTranslator javaTypeTranslator) {
        this.namingPolicy = namingPolicy;
        this.javaTypeTranslator = javaTypeTranslator;
    }

    public TypeDesc getClosure(FunctionValue functionValue, int knownParametersCount) {
        ClosureDesc desc = new ClosureDesc(functionValue, knownParametersCount);
        TypeDesc result = cache.get(desc);
        if(result == null) {
            result = createClosure(functionValue, knownParametersCount);
            cache.put(desc, result);
        }
        return result;
    }

    private TypeDesc createClosure(FunctionValue functionValue, int knownParametersCount) {
        String className = namingPolicy.getFreshClosureClassName();
        
        int remainingArity = functionValue.getArity() - knownParametersCount;
        TypeDesc[] parameterTypes = javaTypeTranslator.toTypeDescs(functionValue.getParameterTypes());
        
        // Create new class
        ClassBuilder classBuilder;
        if(remainingArity <= Constants.MAX_FUNCTION_PARAMETER_COUNT) {
            if(SCLCompilerConfiguration.TRACE_METHOD_CREATION)
                System.out.println("Create class " + className);
            classBuilder = new ClassBuilder(this, Opcodes.ACC_PUBLIC, className, MethodBuilderBase.getClassName(Constants.FUNCTION_IMPL[remainingArity]));
            classBuilder.setSourceFile("_SCL_Closure");
            
            // Create fields
            CodeBuilderUtils.makeRecord(classBuilder, functionValue.toString(), Opcodes.ACC_PRIVATE | Opcodes.ACC_FINAL, "p",
                    Arrays.copyOf(parameterTypes, knownParametersCount),
                    true);
            
            // Create apply
            {
                MethodBuilder mb = classBuilder.addMethod(Opcodes.ACC_PUBLIC, "apply", TypeDesc.OBJECT, Constants.OBJECTS[remainingArity]);
                Val[] parameters = new Val[functionValue.getArity()];
                for(int i=0;i<knownParametersCount;++i)
                    parameters[i] = new LocalFieldConstant(functionValue.getParameterTypes()[i], "p"+i);
                for(int i=0;i<remainingArity;++i) {
                    Type type = functionValue.getParameterTypes()[knownParametersCount+i];
                    parameters[knownParametersCount+i] = new LocalVariableConstant(type, mb.getParameter(i));
                }
                functionValue.prepare(mb);
                Type returnType = functionValue.applyExact(mb, parameters);
                mb.box(returnType);
                mb.returnValue(TypeDesc.OBJECT);
                mb.finish();
            }        
        }
        else {
            if(SCLCompilerConfiguration.TRACE_METHOD_CREATION)
                System.out.println("Create class " + className);
            classBuilder = new ClassBuilder(this, Opcodes.ACC_PUBLIC, className, MethodBuilderBase.getClassName(Constants.FUNCTION_N_IMPL));
            classBuilder.setSourceFile("_SCL_Closure");
            
            // Create fields
            for(int i=0;i<knownParametersCount;++i)
                classBuilder.addField(Opcodes.ACC_PRIVATE | Opcodes.ACC_FINAL, "p"+i, parameterTypes[i]);
            
            // Create constructor
            {
                MethodBuilderBase mb = classBuilder.addConstructorBase(Opcodes.ACC_PUBLIC, Arrays.copyOf(parameterTypes, knownParametersCount));
                mb.loadThis();
                mb.loadConstant(remainingArity);
                mb.invokeConstructor(MethodBuilderBase.getClassName(Constants.FUNCTION_N_IMPL), new TypeDesc[] { TypeDesc.INT });
                for(int i=0;i<knownParametersCount;++i) {
                    mb.loadThis();
                    mb.loadLocal(mb.getParameter(i));
                    mb.storeField(className, "p"+i, parameterTypes[i]);
                }
                mb.returnVoid();
                mb.finish();
            }
            
            // Create apply
            {
                MethodBuilder mb = classBuilder.addMethod(Opcodes.ACC_PUBLIC, "doApply", TypeDesc.OBJECT, new TypeDesc[] {TypeDesc.forClass(Object[].class)});
                Val[] parameters = new Val[functionValue.getArity()];
                for(int i=0;i<knownParametersCount;++i)
                    parameters[i] = new LocalFieldConstant(functionValue.getParameterTypes()[i], "p"+i);
                LocalVariable parameter = mb.getParameter(0);
                for(int i=0;i<remainingArity;++i) {
                    mb.loadLocal(parameter);
                    mb.loadConstant(i);
                    mb.loadFromArray(TypeDesc.OBJECT);
                    Type type = functionValue.getParameterTypes()[knownParametersCount+i];
                    TypeDesc typeDesc = javaTypeTranslator.toTypeDesc(type);
                    mb.unbox(type);
                    LocalVariable lv = mb.createLocalVariable("p"+(i+knownParametersCount), typeDesc);
                    mb.storeLocal(lv);
                    BoundVar var = new BoundVar(type);
                    parameters[knownParametersCount+i] = var;
                    mb.setLocalVariable(var, lv);
                }
                functionValue.applyExact(mb, parameters);
                mb.box(functionValue.getReturnType());
                mb.returnValue(TypeDesc.OBJECT);
            }
            
            CodeBuilderUtils.implementHashCodeAndEquals(classBuilder, functionValue.toString(), "p", parameterTypes);
        }
            
        // Finish
        addClass(classBuilder);
        
        return TypeDesc.forClass(className);
    }
    
    public Map<String, byte[]> getClasses() {
        return classes;
    }
    
    public JavaNamingPolicy getNamingPolicy() {
        return namingPolicy;
    }
    
    public MethodSizeCounter getMethodSizeCounter() {
        return methodSizeCounter;
    }
}