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 com.android.dx.dex.file.DexFile;

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.BufferedOutputStream;
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;

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();
        DexFile dexFile = new 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);
            dexFile.add(CfTranslator.translate(cf, bytes, cfOptions, dexOptions, dexFile));
        }
        String name = randomString.nextString();
        File jar = new File(cacheDir, name + ".jar");
        JarOutputStream out = null;
        dalvik.system.DexFile loadedFile = null;
        try {
            out = new JarOutputStream(new BufferedOutputStream(new FileOutputStream(jar)));
            out.putNextEntry(new JarEntry("classes.dex"));
            dexFile.writeTo(out, null, false);
            out.close();
            loadedFile = dexFileLoader.loadDex(jar, new File(cacheDir, name + ".dex"), 0);
            Map<TypeDescription, Class<?>> classMap = new HashMap<>();
            for (TypeDescription td : types.keySet()) {
                classMap.put(td, loadClass(loadedFile, td, classLoader));
            }
            return classMap;
        } catch (IOException | ClassNotFoundException e) {
            throw new RuntimeException(e);
        } finally {
            if (loadedFile != null) {
                try {
                    loadedFile.close();
                } catch (IOException ignored) {
                }
            }
            IOUtil.closeQuietly(out);
            if (jar.exists() && !jar.delete()) {
                Logger.getLogger("com.github.tmurakami.dexmockito").warning("Cannot delete " + jar);
            }
        }
    }

    private static Class<?> loadClass(dalvik.system.DexFile dexFile,
                                      TypeDescription typeDescription,
                                      ClassLoader classLoader) throws ClassNotFoundException {
        Class<?> c = dexFile.loadClass(typeDescription.getName(), classLoader);
        if (c == null) {
            throw new ClassNotFoundException(typeDescription.getName());
        }
        return c;
    }

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

}

@@ -261,48 +261,8 @@
 tor;
-%0Aimport com.android.dx.dex.file.DexFile;
 %0A%0Aim
@@ -466,45 +466,8 @@
 l;%0A%0A
-import java.io.BufferedOutputStream;%0A
 impo
@@ -695,16 +695,47 @@
 ogger;%0A%0A
+import dalvik.system.DexFile;%0A%0A
 final cl
@@ -1379,24 +1379,48 @@
         
+com.android.dx.dex.file.
 DexFile 
 dexFile 
@@ -1411,20 +1411,17 @@
 DexFile 
-dexF
+f
 ile = ne
@@ -1422,16 +1422,40 @@
 e = new 
+com.android.dx.dex.file.
 DexFile(
@@ -1816,28 +1816,25 @@
             
-dexF
+f
 ile.add(CfTr
@@ -1886,20 +1886,17 @@
 ptions, 
-dexF
+f
 ile));%0A 
@@ -2011,93 +2011,8 @@
 %22);%0A
-        JarOutputStream out = null;%0A        dalvik.system.DexFile loadedFile = null;%0A
     
@@ -2037,18 +2037,8 @@
     
-out = new 
 JarO
@@ -2052,21 +2052,22 @@
 ream
-(new Buffered
+ out = new Jar
 Outp
@@ -2101,19 +2101,40 @@
 am(jar))
-)
 ;%0A
+            try %7B%0A    
         
@@ -2196,20 +2196,21 @@
         
-dexF
+    f
 ile.writ
@@ -2248,39 +2248,170 @@
     
-out.close();%0A            loaded
+%7D finally %7B%0A                IOUtil.closeQuietly(out);%0A            %7D%0A            Map%3CTypeDescription, Class%3C?%3E%3E classMap = new HashMap%3C%3E();%0A            DexFile dex
 File
@@ -2495,67 +2495,18 @@
     
-Map%3CTypeDescription, Class%3C?%3E%3E classMap = new HashMap%3C%3E();%0A
+try %7B%0A    
     
@@ -2573,16 +2573,20 @@
         
+    
 classMap
@@ -2604,22 +2604,19 @@
 adClass(
-loa
 de
-d
+x
 File, td
@@ -2648,196 +2648,35 @@
     
-%7D%0A            return classMap;%0A        %7D catch (IOException %7C ClassNotFoundException e) %7B%0A            throw new RuntimeException(e);%0A        %7D finally %7B%0A            if (loadedFile != null)
+    %7D%0A            %7D finally
  %7B%0A 
@@ -2720,14 +2720,11 @@
     
-loa
 de
-d
+x
 File
@@ -2829,33 +2829,146 @@
     
-IOUtil.closeQuietly(out);
+return classMap;%0A        %7D catch (IOException %7C ClassNotFoundException e) %7B%0A            throw new RuntimeException(e);%0A        %7D finally %7B
 %0A   
@@ -3183,30 +3183,16 @@
 adClass(
-dalvik.system.
 DexFile 

Remove redundant buffering