package me.yamakaja.runtimetransformer.transform;

import me.yamakaja.runtimetransformer.agent.AgentJob;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.ClassNode;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.List;
import java.util.function.IntFunction;
import java.util.stream.Collectors;

/**
 * Created by Yamakaja on 18.05.17.
 */
public class ClassTransformer implements ClassFileTransformer {

    private List<AgentJob> agentJobs;
    private List<Class<?>> classesToRedefine;

    public ClassTransformer(List<AgentJob> agentJobs) {
        this.agentJobs = agentJobs;
        classesToRedefine = agentJobs.stream().map(AgentJob::getToTransform).collect(Collectors.toList());
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        if (!classesToRedefine.contains(classBeingRedefined))
            return classfileBuffer;

        ClassWriter writer;
        try {
            ClassReader reader = new ClassReader(classfileBuffer);

            ClassNode node = new ClassNode(Opcodes.ASM5);
            reader.accept(node, 0);

            this.agentJobs.stream()
                    .filter(job -> job.getToTransform().getName().replace('.', '/').equals(className))
                    .forEach(job -> job.apply(node));

            writer = new FixedClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES, loader);

            node.accept(writer);
        } catch (Throwable e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }

        return writer.toByteArray();
    }

    public Class<?>[] getClassesToTransform() {
        return agentJobs.stream().map(AgentJob::getToTransform).toArray((IntFunction<Class<?>[]>) Class[]::new);
    }

}