package me.coley.recaf.workspace;

import me.coley.recaf.control.Controller;
import me.coley.recaf.util.ClasspathUtil;
import me.coley.recaf.util.IOUtil;
import me.coley.recaf.util.Log;
import org.objectweb.asm.Type;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.lang.module.Module;
import java.security.ProtectionDomain;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Importable instrumentation resource.
 *
 * @author Matt
 */
public class InstrumentationResource extends JavaResource {
	private static final ResourceLocation LOCATION = LiteralResourceLocation.ofKind(
			ResourceKind.INSTRUMENTATION,
			"Instrumentation");
	public static Instrumentation instrumentation;
	private static InstrumentationResource instance;

	/**
	 * Constructs an instrumentation resource.
	 *
	 * @throws IllegalStateException
	 * 		When the {@link #instrumentation} instance has not been set.
	 * @throws IOException
	 * 		When querying for runtime classes fails.
	 */
	private InstrumentationResource() throws IllegalStateException, IOException {
		super(ResourceKind.INSTRUMENTATION);
		// Instrumentation is ALWAYS primary
		setPrimary(true);
		if(instrumentation == null)
			throw new IllegalStateException("Instrumentation has not been initialized!");
		if (instance != null)
			throw new IllegalStateException("There already is an instrumentation resource!");
		instance = this;
		setSkippedPrefixes(Arrays.asList("java/", "javax/", "javafx/", "sun/",
				"com/sun/", "com/oracle/", "jdk/", "me/coley/"));
	}

	/**
	 * Setup an instrumentation based workspace.
	 *
	 * @param controller Controller to act on.
	 *
	 * @return Created workspace.
	 */
	public static Workspace setup(Controller controller) {
		try {
			// Add transformer to add new classes to the map
			instrumentation.addTransformer(new InstrumentationResourceTransformer());
			// Set workspace
			controller.setWorkspace(new Workspace(getInstance()));
			Log.info("Loaded instrumentation workspace");
		} catch(Exception ex) {
			Log.error(ex, "Failed to initialize instrumentation");
		}
		return controller.getWorkspace();
	}

	/**
	 * Saves changed by retransforming classes.
	 *
	 * @throws ClassNotFoundException
	 * 		When the modified class couldn't be found.
	 * @throws UnmodifiableClassException
	 * 		When the modified class is not allowed to be modified.
	 * @throws ClassFormatError
	 * 		When the modified class is not valid.
	 */
	public void save() throws ClassNotFoundException, UnmodifiableClassException, ClassFormatError {
		// Classes to update
		Set<String> dirty = new HashSet<>(getDirtyClasses());
		if(dirty.isEmpty()) {
			Log.info("There are no classes to redefine.", dirty.size());
			return;
		}
		Log.info("Preparing to redefine {} classes", dirty.size());
		ClassDefinition[] definitions = new ClassDefinition[dirty.size()];
		int i = 0;
		for (String name : dirty) {
			String clsName = name.replace('/', '.');
			Class<?> cls = Class.forName(clsName, false, ClasspathUtil.scl);
			byte[] value = getClasses().get(name);
			if (value == null)
				throw new IllegalStateException("Failed to fetch code for class: " + name);
			definitions[i] = new ClassDefinition(cls, value);
			i++;
		}
		// Apply new definitions
		instrumentation.redefineClasses(definitions);
		// We don't want to continually re-apply changes that don't need to be updated
		getDirtyClasses().clear();
		Log.info("Successfully redefined {} classes", definitions.length);
	}

	@Override
	protected Map<String, byte[]> loadClasses() throws IOException {
		return Collections.emptyMap();
	}

	@Override
	protected Map<String, byte[]> loadFiles() {
		return Collections.emptyMap();
	}

	@Override
	protected Map<String, byte[]> copyMap(Map<String, byte[]> map) {
		return new ConcurrentHashMap<>(map);
	}

	@Override
	public String toString() {
		return "Instrumentation";
	}

	private void loadRuntimeClasses(Map<String, byte[]> map) throws IOException {
		// iterate over loaded classes
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		byte[] buffer = new byte[8192];
		for(Class<?> c : instrumentation.getAllLoadedClasses()) {
			String name = Type.getInternalName(c);
			// skip specified prefixes
			if(shouldSkip(name))
				continue;
			// Skip array types
			if (name.contains("["))
				continue;
			String path = name.concat(".class");
			ClassLoader loader = c.getClassLoader();
			try(InputStream in = (loader != null) ?
					loader.getResourceAsStream(path) :
					ClassLoader.getSystemResourceAsStream(path)) {
				if(in != null) {
					out.reset();
					getClasses().put(name, IOUtil.toByteArray(in, out, buffer));
					getDirtyClasses().remove(name);
				}
			}
		}
	}

	/**
	 * @return Instrumentation resource instance.
	 *
	 * @throws IOException
	 * 		When the resource cannot be instantiated.
	 */
	public static InstrumentationResource getInstance() throws IOException {
		if (instance == null)
			instance = new InstrumentationResource();
		return instance;
	}

	/**
	 * @return {@code true} if Recaf is running as a Java agent.
	 */
	public static boolean isActive() {
		return instrumentation != null;
	}

	/**
	 * Transformer to load classes from instrumentation.
	 */
	private static class InstrumentationResourceTransformer implements ClassFileTransformer {
		private boolean firstTransformerLoad = true;

		public byte[] transform(Module module, ClassLoader loader, String className,
								Class<?> cls, ProtectionDomain domain, byte[] buffer) {
			return transform(loader, className, cls, domain, buffer);
		}

		@Override
		public byte[] transform(ClassLoader loader, String className,
								Class<?> cls, ProtectionDomain domain, byte[] buffer) {
			// This super odd way of getting the resource IS INTENTIONAL.
			// If you choose to optimize this in the future verify it behaves the same.
			InstrumentationResource res = null;
			try {
				res = getInstance();
				synchronized (this) {
					if (firstTransformerLoad) {
						firstTransformerLoad = false;
						// There is a time gap between when we first called 'loadClasses' and this gets called.
						// We need to fetch those classes here so we have everything available.
						res.loadRuntimeClasses(getInstance().getClasses());
					}
				}
			} catch(IOException ex) { return buffer; }
			// Check to skip class
			String internal = className.replace('.', '/');
			if(res.shouldSkip(internal))
				return buffer;
			// Add to classes map
			res.getClasses().put(internal, buffer);
			// Make sure the class is NOT marked as dirty after initially registering it
			res.getDirtyClasses().remove(internal);
			return buffer;
		}
	}

	@Override
	public ResourceLocation getShortName() {
		return LOCATION;
	}

	@Override
	public ResourceLocation getName() {
		return LOCATION;
	}
}