package com.hammurapi.common;

import java.io.File;
import java.io.FileInputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.jar.Attributes;
import java.util.jar.Manifest;
import java.util.logging.Logger;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

import org.apache.bcel.classfile.Constant;
import org.apache.bcel.classfile.ConstantClass;
import org.apache.bcel.classfile.ConstantPool;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.util.ClassLoaderRepository;

/**
 * Scans class dependencies and returns a list of classpath entries URL's which given class depends on.
 * This class requires BCEL 5.2 in the classpath.
 * @author Pavel Vlasov
 *
 */
public class DependencyScanner {
	
	private static final String CLASS_EXTENSION = ".class";

	private static final Logger logger = Logger.getLogger(DependencyScanner.class.getName());
	
	private static Map<ClassLoader, ClassLoaderRepository> repoCache = new java.util.HashMap<ClassLoader, ClassLoaderRepository>();
	
	private static URL classContainer(Class<?> clazz) {
		String className = clazz.getName();
		int lastDotIdx = className.lastIndexOf('.');
		String shortName = lastDotIdx==-1 ? clazz.getName() : className.substring(lastDotIdx+1);
		URL classResource = clazz.getResource(shortName+CLASS_EXTENSION);
		if (classResource!=null) {
			if ("jar".equalsIgnoreCase(classResource.getProtocol())) {
				String path = classResource.getPath();				
				int exlaIdx = path.indexOf('!');
				if (exlaIdx!=-1) {
					path = path.substring(0, exlaIdx);
					try {
						return new URL(path);
					} catch (MalformedURLException e) {
						logger.warning(e.toString());
					}
				}
			} else {
				String postfix = clazz.getName().replace('.', '/')+CLASS_EXTENSION;
				String cref = classResource.toExternalForm();
				if (cref.endsWith(postfix)) {
					try {
						return new URL(cref.substring(0, cref.length()-postfix.length()-1));
					} catch (MalformedURLException e) {
						logger.warning(e.toString());
					}					
				}
			}
		}
		return null;
	}

	public static List<URL> scan(Class<?> clazz) {		
		List<URL> ret = new ArrayList<URL>();
		if (clazz.getClassLoader()!=null) {
			URL cc = classContainer(clazz);
			if (cc!=null) {
				ret.add(cc);
				URL rtJar = classContainer(Object.class);
				String rtJarPath = rtJar.getPath();
				int idx = rtJarPath.lastIndexOf('/');
				scanEntry(ret, clazz.getClassLoader(), 0, idx == -1 ? rtJarPath : rtJarPath.substring(0, idx));
			}
		}
		return ret;
	}
	
	private static void scanEntry(List<URL> collector, ClassLoader classLoader, int idx, String jreLibDir) {
		URL toScan = collector.get(idx);
		String toScanPath = toScan.getPath();
		int slashIdx = toScanPath.lastIndexOf('/');
		if (slashIdx!=-1 && toScanPath.substring(0, slashIdx).equals(jreLibDir)) {
			scanEntry(collector, classLoader, idx+1, jreLibDir);
			return;
		}
		
		ClassLoaderRepository repo = repoCache.get(classLoader);
		if (repo == null) {
			repo = new ClassLoaderRepository(classLoader);
			repoCache.put(classLoader, repo);
		}
		
		try {
			if ("file".equalsIgnoreCase(toScan.getProtocol()) && new File(toScan.toURI()).isDirectory()) {
				for (File child: new File(toScan.toURI()).listFiles()) {
					scanFile(collector, child, null, repo, classLoader, jreLibDir);
				}
			} else {
				ZipInputStream zis = new ZipInputStream(toScan.openStream());
				ZipEntry ze;
				while ((ze = zis.getNextEntry())!=null) {
					if ("META-INF/manifest.mf".equalsIgnoreCase(ze.getName())) {
						Manifest manifest = new Manifest(zis);
						Attributes mattrs = manifest.getMainAttributes();
						String classPath = mattrs.getValue("Class-Path");
						if (classPath!=null) {
							for (String cpe: classPath.split(" ")) {
								try {
									URL reference = new URL(toScan, cpe);
									if (!collector.contains(reference)) {
										collector.add(reference);
									}
								} catch (MalformedURLException mee) {
									logger.warning(mee.toString());
								}
							}
						}
					} else if (ze.getName().endsWith(CLASS_EXTENSION)) {
						JavaClass jc = repo.loadClass(ze.getName().substring(0, ze.getName().length() - CLASS_EXTENSION.length()).replace('/', '.'));
						ConstantPool constantPool = jc.getConstantPool();
						for (Constant constant: constantPool.getConstantPool()) {
							if (constant instanceof ConstantClass) {
								ConstantClass cc = (ConstantClass) constant;
								String cName = String.valueOf(cc.getConstantValue(constantPool)).replace('/', '.');
								if (cName.startsWith("[")) {
									continue;
								}
								URL container = classContainer(classLoader.loadClass(cName));
								if (!collector.contains(container)) {
									String containerPath = container.getPath();
									int sIdx = containerPath.lastIndexOf('/');
									if (sIdx!=-1 && containerPath.substring(0, sIdx).equals(jreLibDir)) {
										continue;
									}
									collector.add(container);
								}
							}
						}
						
					}
	 			}
			}			
		} catch (Exception e) {
			logger.warning(e.toString());
		}
		
	}
	
	private static void scanFile(List<URL> collector, File file, String path, ClassLoaderRepository repo, ClassLoader classLoader, String jreLibDir) throws Exception {
		String fullName = path==null ? file.getName() : path+"/"+file.getName();
		if (file.isDirectory()) {
			for (File child: file.listFiles()) {
				scanFile(collector, child, fullName, repo, classLoader, jreLibDir);
			}
		} else if ("META-INF/manifest.mf".equalsIgnoreCase(fullName)) {
			Manifest manifest = new Manifest(new FileInputStream(file));
			Attributes mattrs = manifest.getMainAttributes();
			String classPath = mattrs.getValue("Class-Path");
			if (classPath!=null) {
				for (String cpe: classPath.split(" ")) {
					try {
						URL reference = new URL(file.toURI().toURL(), cpe);
						if (!collector.contains(reference)) {
							collector.add(reference);
						}
					} catch (MalformedURLException mee) {
						logger.warning(mee.toString());
					}
				}
			}
		} else if (file.getName().endsWith(CLASS_EXTENSION)) {
			JavaClass jc = repo.loadClass(fullName.substring(0, fullName.length() - CLASS_EXTENSION.length()).replace('/', '.'));
			ConstantPool constantPool = jc.getConstantPool();
			for (Constant constant: constantPool.getConstantPool()) {
				if (constant instanceof ConstantClass) {
					ConstantClass cc = (ConstantClass) constant;
					String cName = String.valueOf(cc.getConstantValue(constantPool)).replace('/', '.');
					if (cName.startsWith("[")) {
						continue;
					}
					URL container = classContainer(classLoader.loadClass(cName));
					if (!collector.contains(container)) {
						String containerPath = container.getPath();
						int sIdx = containerPath.lastIndexOf('/');
						if (sIdx!=-1 && containerPath.substring(0, sIdx).equals(jreLibDir)) {
							continue;
						}
						collector.add(container);
					}
				}
			}
			
		}
	}

	public static void main(String[] args) throws Exception {

		List<URL> classPath = scan(DependencyScanner.class);
		for (URL entry: classPath) {
			System.out.println(entry);
		}
		
	}

}
