package org.stagemonitor.requestmonitor.ejb;

import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
import static net.bytebuddy.matcher.ElementMatchers.isPublic;

import javax.ejb.Remote;

import net.bytebuddy.description.annotation.AnnotationList;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import org.stagemonitor.core.util.ClassUtils;
import org.stagemonitor.requestmonitor.AbstractMonitorRequestsTransformer;

public class RemoteEjbMonitorTransformer extends AbstractMonitorRequestsTransformer {

	@Override
	protected ElementMatcher.Junction<TypeDescription> getIncludeTypeMatcher() {
		return super.getIncludeTypeMatcher().and(isAnnotatedWith(Remote.class));
	}

	@Override
	protected ElementMatcher.Junction<? super MethodDescription.InDefinedShape> getExtraMethodElementMatcher() {
		return isPublic().and(new IsDeclaredInRemoteClassElementMatcher());
	}

	@Override
	public boolean isActive() {
		return ClassUtils.isPresent("javax.ejb.Remote");
	}

	private static class IsDeclaredInRemoteClassElementMatcher extends ElementMatcher.Junction.AbstractBase<MethodDescription.InDefinedShape> {
		@Override
		public boolean matches(MethodDescription.InDefinedShape targetMethod) {
			final AnnotationList declaredAnnotationsOfType = targetMethod.getDeclaringType().getDeclaredAnnotations();
			if (!declaredAnnotationsOfType.isAnnotationPresent(Remote.class)) {
				return false;
			}
			final Remote annotation = declaredAnnotationsOfType.ofType(Remote.class).loadSilent();
			return isDeclaredInAnyRemoteInterface(targetMethod, annotation.value());
		}

		private boolean isDeclaredInAnyRemoteInterface(MethodDescription.InDefinedShape targetMethod, Class[] remoteInterfaces) {
			for (Class<?> remoteInterface : remoteInterfaces) {
				if (isDeclaredIn(targetMethod, remoteInterface)) {
					return true;
				}
			}
			return false;
		}

		private boolean isDeclaredIn(final MethodDescription.InDefinedShape targetMethod, Class<?> remoteInterface) {
			final TypeDescription.ForLoadedType remoteInterfaceType = new TypeDescription.ForLoadedType(remoteInterface);
			return remoteInterfaceType
					.getDeclaredMethods()
					.filter(new ElementMatcher<MethodDescription.InDefinedShape>() {
				@Override
				public boolean matches(MethodDescription.InDefinedShape target) {
					return target.getName().equals(targetMethod.getName()) &&
							target.getDescriptor().equals(targetMethod.getDescriptor());
				}
			}).size() > 0;
		}
	}
}

@@ -169,16 +169,204 @@
 sPublic;
+%0Aimport static net.bytebuddy.matcher.ElementMatchers.named;%0Aimport static net.bytebuddy.matcher.ElementMatchers.returns;%0Aimport static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 %0A%0Aimport
@@ -551,32 +551,80 @@
 ypeDescription;%0A
+import net.bytebuddy.description.type.TypeList;%0A
 import net.byteb
@@ -1663,13 +1663,15 @@
 %0A%09%09%09
-if (!
+return 
 decl
@@ -1729,64 +1729,46 @@
 ass)
-) %7B
+ &&
 %0A%09%09%09%09
-return false;%0A%09%09%09%7D%0A%09%09%09final Remote annotation = 
+%09!new TypeList.ForLoadedTypes(
 decl
@@ -1826,83 +1826,87 @@
 nt()
-;%0A%09%09%09return isDeclaredInAnyRemoteInterface(targetMethod, annotation.value()
+.value())%0A%09%09%09%09%09%09%09.filter(new IsDeclaredInElementMatcher(targetMethod)).isEmpty(
 );%0A%09
@@ -1902,32 +1902,35 @@
 sEmpty();%0A%09%09%7D%0A%0A%09
+%7D%0A%0A
 %09private boolean
@@ -1926,637 +1926,284 @@
 ate 
-boolean isDeclaredInAnyRemoteInterface(MethodDescription.InDefinedShape targetMethod, Class%5B%5D remoteInterfaces) %7B%0A%09%09%09for (Class%3C?%3E remoteInterface : remoteInterfaces) %7B%0A%09%09%09%09if (isDeclaredIn(targetMethod, remoteInterface)) %7B%0A%09%09%09%09%09return true;%0A%09%09%09%09%7D%0A%09%09%09%7D%0A%09%09%09return false;%0A%09%09%7D%0A%0A%09%09private boolean isDeclaredIn(final MethodDescription.InDefinedShape targetMethod, Class%3C?%3E remoteInterface) %7B%0A%09%09%09final TypeDescription.ForLoadedType remoteInterfaceType = new TypeDescription.ForLoadedType(remoteInterface);%0A%09%09%09return remoteInterfaceType%0A%09%09%09%09%09.getDeclared
+static class IsDeclaredInElementMatcher implements ElementMatcher%3CTypeDescription%3E %7B%0A%09%09private final MethodDescription.InDefinedShape targetMethod;%0A%0A%09%09public IsDeclaredInElementMatcher(MethodDescription.InDefinedShape target
 Method
-s()
+) %7B
 %0A%09%09%09
-%09%09.filter(new ElementMatcher%3CMethodDescription.InDefinedShape%3E() %7B%0A%09%09
+this.targetMethod = targetMethod;%0A%09%09%7D%0A%0A
 %09%09@O
@@ -2210,18 +2210,16 @@
 verride%0A
-%09%09
 %09%09public
@@ -2227,38 +2227,36 @@
 boolean matches(
-Method
+Type
 Description.InDe
@@ -2246,39 +2246,24 @@
 eDescription
-.InDefinedShape
  target) %7B%0A%09
@@ -2261,18 +2261,16 @@
 rget) %7B%0A
-%09%09
 %09%09%09retur
@@ -2271,16 +2271,17 @@
 %09return 
+!
 target.g
@@ -2286,21 +2286,44 @@
 .get
-Name().equals
+DeclaredMethods()%0A%09%09%09%09%09.filter(named
 (tar
@@ -2346,19 +2346,16 @@
 e())
- &&
 %0A%09%09%09%09%09%09%09
 targ
@@ -2354,91 +2354,164 @@
 %09%09%09%09
-target.getDescriptor().equals(targetMethod.getDescriptor());
+.and(returns(targetMethod.getReturnType().asErasure()))%0A%09%09%09%09%09%09%09.and(takesArguments(targetMethod.getParameters().asTypeList().asErasures())))
 %0A%09%09%09%09
-%7D%0A%09%09%09%7D).size() %3E 0
+%09.isEmpty()
 ;%0A%09%09

Improve remote EJB call matcher