package yushijinhun.authlibagent.javaagent;

import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Collection;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.logging.Logger;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import yushijinhun.authlibagent.javaagent.TransformUnit.TransformCheckCallback;

public class Transformer implements ClassFileTransformer {

	private static final Logger LOGGER = Logger.getLogger("authlibagent.transformer");

	private static class TransformCheckCallbackImpl implements TransformCheckCallback {

		boolean isAllowed = false;

		@Override
		public void allowTransform() {
			isAllowed = true;
		}

	}

	private Map<String, Collection<TransformUnit>> units = new ConcurrentHashMap<>();
	private boolean debugMode;

	protected void addTransformUnit(String className, TransformUnit unit) {
		Objects.requireNonNull(className);
		Objects.requireNonNull(unit);
		synchronized (units) {
			Collection<TransformUnit> classUnits = units.get(className);
			if (classUnits == null) {
				classUnits = new CopyOnWriteArrayList<>();
				units.put(className, classUnits);
			}
			classUnits.add(unit);
		}
	}

	public boolean setDebugMode() {
		return debugMode;
	}

	@Override
	public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
		Collection<TransformUnit> classunits = units.get(className);
		if (classunits != null) {
			LOGGER.info("try to transform " + className);
			byte[] output = transformClass(classunits, classfileBuffer);
			if (output != null) {
				LOGGER.info("transform " + className);
				if (debugMode) {
					debugSaveModifiedClass(output, className);
				}
				return output;
			}
		}
		return null;
	}

	private void debugSaveModifiedClass(byte[] classBuffer, String className) {
		className = className.replace('/', '.');
		try (OutputStream out = new FileOutputStream(className + "_modified.class")) {
			out.write(classBuffer);
		} catch (IOException e) {
			LOGGER.warning("failed to save modified class " + className + "　: " + e);
			e.printStackTrace();
		}
	}

	private byte[] transformClass(Collection<TransformUnit> classunits, byte[] classbuffer) {
		if (units == null) {
			return null;
		}

		boolean changed = false;
		byte[] currentClass = classbuffer;
		ClassReader classreader = new ClassReader(currentClass);
		for (TransformUnit unit : classunits) {
			TransformCheckCallbackImpl callback = new TransformCheckCallbackImpl();
			ClassVisitor checker = unit.check(callback);
			if (checker != null) {
				classreader.accept(checker, ClassReader.SKIP_DEBUG);
			}
			if (callback.isAllowed) {
				ClassWriter classwriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
				ClassVisitor modifier = unit.transform(classwriter);
				classreader.accept(modifier, ClassReader.SKIP_DEBUG);
				currentClass = classwriter.toByteArray();
				classreader = new ClassReader(currentClass);
				changed = true;
			}
		}
		if (changed) {
			return currentClass;
		}
		return null;
	}

}

@@ -1267,16 +1267,59 @@
 (unit);%0A
+%09%09className = className.replace('.', '/');%0A
 %09%09synchr
@@ -1808,16 +1808,25 @@
 ption %7B%0A
+%09%09try %7B%0A%09
 %09%09Collec
@@ -1880,16 +1880,17 @@
 sName);%0A
+%09
 %09%09if (cl
@@ -1908,24 +1908,25 @@
  null) %7B%0A%09%09%09
+%09
 LOGGER.info(
@@ -1959,16 +1959,17 @@
 sName);%0A
+%09
 %09%09%09byte%5B
@@ -2023,24 +2023,25 @@
 Buffer);%0A%09%09%09
+%09
 if (output !
@@ -2042,32 +2042,33 @@
 tput != null) %7B%0A
+%09
 %09%09%09%09LOGGER.info(
@@ -2094,24 +2094,25 @@
 sName);%0A%09%09%09%09
+%09
 if (debugMod
@@ -2112,24 +2112,25 @@
 ebugMode) %7B%0A
+%09
 %09%09%09%09%09debugSa
@@ -2169,18 +2169,20 @@
 e);%0A%09%09%09%09
+%09
 %7D%0A
+%09
 %09%09%09%09retu
@@ -2191,25 +2191,133 @@
  output;%0A%09%09%09
-%7D
+%09%7D%0A%09%09%09%7D%0A%09%09%7D catch (Throwable e) %7B%0A%09%09%09LOGGER.info(%22failed to transform %22 + className);%0A%09%09%09e.printStackTrace();
 %0A%09%09%7D%0A%09%09retur
@@ -3229,24 +3229,67 @@
 sAllowed) %7B%0A
+%09%09%09%09LOGGER.info(%22transform unit %22 + unit);%0A
 %09%09%09%09ClassWri

add logging and use canonical class name