package kindle.test.jni;

import java.io.BufferedInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.net.JarURLConnection;
import java.net.URL;
import java.security.CodeSource;
import java.security.SecureClassLoader;
import java.security.cert.Certificate;
import java.util.HashMap;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

public class JniClassLoader extends SecureClassLoader {

	public class ClassData {
		   protected byte[] data;
		   protected Certificate[] certs;
		   protected URL codebase;

		   public ClassData(byte[] data, Certificate[] certs, URL codebase) {
			   this.data = data;
			   this.certs  = certs;
			   this.codebase = codebase;
		   }
		   public byte[] getData() {
		      return data;
		   }
		   public Certificate[] getCerts() {
		      return certs;
		   }
		   public URL getCodeBase() {
			      return codebase;
		   }
	}

    private Map<String, Class<?>> classes;

    public JniClassLoader() {
        super(JniClassLoader.class.getClassLoader());
        classes = new HashMap<String, Class<?>>();
    }

    public String toString() {
        return JniClassLoader.class.getName();
    }

    @Override
    public Class<?> findClass(String name) throws ClassNotFoundException {

        if (classes.containsKey(name)) {
            return classes.get(name);
        }

        String path = name.replace('.', File.separatorChar) + ".class";
        ClassData cls = null;

        try {
        	cls = loadClassData(path);
        } catch (IOException e) {
            throw new ClassNotFoundException("Class not found at path: " + name, e);
        }

		//Unfortunately we cannot use certs read from azw2 file, because they causes problem with SecurityManager.
		//They are not recognized by policy with option "signBy Kindle KindleInteractive etc.".
		//If someone know why please let me know.
		//Here I will use certs copied from class that was created by default classloader - they are ok. 
		CodeSource cs = new CodeSource(cls.getCodeBase(), (Certificate[]) JniClassLoader.class.getSigners());

        Class<?> c = defineClass(name, cls.getData(), 0, cls.getData().length, cs);

        resolveClass(c);
        classes.put(name, c);

        return c;
    }

    /**
     * @param name - path of the class in azw2 file
     * @return
     * @throws IOException
     * 
     *	http://docstore.mik.ua/orelly/java-ent/security/ch12_02.htm
     */
    private ClassData loadClassData(String name) throws IOException {
        byte[] classBytes = null;
        Certificate certs[] = new Certificate[0];
        URL codebase = null;

    	URL url = Thread.currentThread().getContextClassLoader().getResource(name);
		
        JarURLConnection jarConn = (JarURLConnection) url.openConnection();
        codebase = jarConn.getJarFileURL();
        
        JarInputStream jis = new JarInputStream(codebase.openConnection().getInputStream());
		JarEntry je = null;
        
		while ((je = jis.getNextJarEntry()) != null) {
			if (name.equals(je.getName())) {
				BufferedInputStream jarBuf = new BufferedInputStream(jis);
				ByteArrayOutputStream jarOut = new ByteArrayOutputStream();

				int b;
				while ((b = jarBuf.read()) != -1)
					jarOut.write(b);

				certs = je.getCertificates();
				if (null == certs)
					certs = new Certificate[0];

				classBytes = jarOut.toByteArray();
				jis.closeEntry();
				break;
			}
			jis.closeEntry();
		}
		
		if (null == classBytes)
			throw new IOException("Failed to load "+name+" class: "+url);

        return new ClassData(classBytes, certs, codebase);
    }

    public void finalize() {
    }
}