package com.github.tmurakami.dexmockito;

import com.android.dx.cf.direct.DirectClassFile;
import com.android.dx.cf.direct.StdAttributeFactory;
import com.android.dx.dex.DexOptions;
import com.android.dx.dex.cf.CfOptions;
import com.android.dx.dex.cf.CfTranslator;

import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
import net.bytebuddy.utility.RandomString;

import org.mockito.internal.util.io.IOUtil;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.logging.Logger;

import dalvik.system.DexFile;

final class DexClassLoadingStrategy implements ClassLoadingStrategy {

    private final File cacheDir;
    private final DexFileLoader dexFileLoader;
    private final RandomString randomString = new RandomString();

    DexClassLoadingStrategy(File cacheDir, DexFileLoader dexFileLoader) {
        this.cacheDir = cacheDir;
        this.dexFileLoader = dexFileLoader;
    }

    @Override
    public Map<TypeDescription, Class<?>> load(ClassLoader classLoader,
                                               Map<TypeDescription, byte[]> types) {
        DexOptions dexOptions = new DexOptions();
        CfOptions cfOptions = new CfOptions();
        com.android.dx.dex.file.DexFile file = new com.android.dx.dex.file.DexFile(dexOptions);
        for (Map.Entry<TypeDescription, byte[]> entry : types.entrySet()) {
            String path = entry.getKey().getName().replace('.', '/') + ".class";
            byte[] bytes = entry.getValue();
            DirectClassFile cf = new DirectClassFile(bytes, path, false);
            cf.setAttributeFactory(StdAttributeFactory.THE_ONE);
            file.add(CfTranslator.translate(cf, bytes, cfOptions, dexOptions, file));
        }
        String name = randomString.nextString();
        File jar = new File(cacheDir, name + ".jar");
        DexFile dexFile;
        try {
            JarOutputStream out = new JarOutputStream(new FileOutputStream(jar));
            try {
                out.putNextEntry(new JarEntry("classes.dex"));
                file.writeTo(out, null, false);
            } finally {
                IOUtil.closeQuietly(out);
            }
            dexFile = dexFileLoader.loadDex(jar, new File(cacheDir, name + ".dex"), 0);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            if (jar.exists() && !jar.delete()) {
                Logger.getLogger("com.github.tmurakami.dexmockito").warning("Cannot delete " + jar);
            }
        }
        Map<TypeDescription, Class<?>> classMap = new HashMap<>();
        try {
            for (TypeDescription td : types.keySet()) {
                Class<?> c = dexFile.loadClass(td.getName(), classLoader);
                if (c == null) {
                    throw new RuntimeException("Cannot load " + td.getName());
                }
                classMap.put(td, c);
            }
        } finally {
            try {
                dexFile.close();
            } catch (IOException ignored) {
            }
        }
        return classMap;
    }

    interface DexFileLoader {
        dalvik.system.DexFile loadDex(File sourceFile, File outputFile, int flags) throws IOException;
    }

}

@@ -1919,17 +1919,21 @@
  String 
-n
+fileN
 ame = ra
@@ -1987,33 +1987,37 @@
  File(cacheDir, 
-n
+fileN
 ame + %22.jar%22);%0A 
@@ -2038,16 +2038,23 @@
  dexFile
+ = null
 ;%0A      
@@ -2420,17 +2420,21 @@
 cheDir, 
-n
+fileN
 ame + %22.
@@ -2452,275 +2452,8 @@
     
-    %7D catch (IOException e) %7B%0A            throw new RuntimeException(e);%0A        %7D finally %7B%0A            if (jar.exists() && !jar.delete()) %7B%0A                Logger.getLogger(%22com.github.tmurakami.dexmockito%22).warning(%22Cannot delete %22 + jar);%0A            %7D%0A        %7D%0A
     
@@ -2515,30 +2515,16 @@
 ap%3C%3E();%0A
-        try %7B%0A
         
@@ -2591,20 +2591,51 @@
     
-Class%3C?%3E c =
+String name = td.getName();%0A               
  dex
@@ -2649,28 +2649,20 @@
 adClass(
-td.getN
+n
 ame
-()
 , classL
@@ -2689,24 +2689,66 @@
     
-if (c == null) %7B
+classMap.put(td, Class.forName(name, false, classLoader));
 %0A   
@@ -2756,24 +2756,26 @@
         
+%7D%0A
         
 throw ne
@@ -2770,153 +2770,183 @@
     
-throw new RuntimeException(%22Cannot load %22 + td.getName());%0A                %7D%0A                classMap.put(td, c);%0A            %7D%0A        %7D finally
+    return classMap;%0A        %7D catch (IOException %7C ClassNotFoundException e) %7B%0A            throw new RuntimeException(e);%0A        %7D finally %7B%0A            if (dexFile != null)
  %7B%0A 
@@ -2948,38 +2948,46 @@
 ) %7B%0A            
+    
 try %7B%0A
+    
                 
@@ -3007,32 +3007,36 @@
 ();%0A            
+    
 %7D catch (IOExcep
@@ -3055,34 +3055,42 @@
 ) %7B%0A            
-%7D%0A
+    %7D%0A    
         %7D%0A      
@@ -3091,32 +3091,181 @@
         
-return classMap;
+    if (jar.exists() && !jar.delete()) %7B%0A                Logger.getLogger(%22com.github.tmurakami.dexmockito%22).warning(%22Cannot delete %22 + jar);%0A            %7D%0A        %7D
 %0A    %7D%0A%0A

Use Class#forName to load classes