 package nxm.sys.lib;

import java.lang.annotation.Documented;
import java.lang.annotation.Inherited;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Function;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Stream;

import nxm.sys.inc.SpecialLoggingRequired;

/** Methods related to inspection of the call stack. These methods generally
 *  apply to logging-related functions (which use the call-stack to identify
 *  the caller), but can be used for other purposes.
 *  @author J. Czechowski
 *  @since  NeXtMidas 3.9.0
 */
@SpecialLoggingRequired
@SuppressWarnings({"BroadCatchBlock", "TooBroadCatch", "UseSpecificCatch"})
public final class CallStackUtil {

  /** prevent instantiation */
  private CallStackUtil () { }

  /** We need to use the traditional logger methods from this class since the methods in this class
   *  have special initialization and calling restrictions due to their use by LogUtilities
   */
  private static final Logger LOG = Logger.getLogger(CallStackUtil.class.getName());

  private static final Object javaLangAccess;
  private static final Method getStackTraceDepth;
  private static final Method getStackTraceElement;
  private static final Object stackWalker;
  private static final Method walk;
  private static final Method toStackTraceElement;
  private static final Method getClassNameStackWalker;

  static {
    boolean currentJavaIs8 = Shell.getJavaMajorVersion() == 8;
    if (!currentJavaIs8 && CallStackUtilJDKSwap.JDK8_CALLSTACKUTIL) {
      LOG.log(Level.INFO, "To get fast implementation, recompile with current version of Java ("+Shell.getJavaMajorVersion()+") or run with JDK 8.\n"+
                          "This code was originally compiled with JDK 8 and does not include the fast version for JDK9+");
    }
    if (currentJavaIs8) {
      // The methods used in getCallerFast(..) are non-standard but provide a measurable performance
      // benefit. As such we get them here and then use then later.
      Object _javaLangAccess;
      Method _getStackTraceDepth;
      Method _getStackTraceElement;
      try {
        Class<?> SharedSecrets        = Class.forName("sun.misc.SharedSecrets");
        Class<?> JavaLangAccess       = Class.forName("sun.misc.JavaLangAccess");
        Method   getJavaLangAccess    = SharedSecrets.getMethod("getJavaLangAccess");

        _javaLangAccess       = getJavaLangAccess.invoke(null);
        _getStackTraceDepth   = JavaLangAccess.getMethod("getStackTraceDepth",   Throwable.class);
        _getStackTraceElement = JavaLangAccess.getMethod("getStackTraceElement", Throwable.class, Integer.TYPE);

        // Now try using it (we assume accessing element 0 is OK as that would represent the class
        // initializer that this code is in).
        Throwable         t     = new Throwable();
        int               depth = (Integer)_getStackTraceDepth.invoke(_javaLangAccess, t);
        StackTraceElement elem  = (StackTraceElement)_getStackTraceElement.invoke(_javaLangAccess, t, 0);
      }
      catch (Throwable t) {
        LOG.log(Level.CONFIG, "Current configuration does not support fast implementation of getCaller(..)", t);
        _javaLangAccess       = null;
        _getStackTraceDepth   = null;
        _getStackTraceElement = null;
      }

      javaLangAccess       = _javaLangAccess;
      getStackTraceDepth   = _getStackTraceDepth;
      getStackTraceElement = _getStackTraceElement;

      // Assign JDK9+ fields to null
      stackWalker             = null;
      walk                    = null;
      toStackTraceElement     = null;
      getClassNameStackWalker = null;
    }
    else {
      // If we can't use the jdk 1.8 way, lets use the new jdk 9 way, the StackWalker API.
      // Because it is not available in jdk 1.8, we need to use reflection for this.
      Method getStackWalker;
      Object _stackWalker;
      Method _walk;
      Method _toStackTraceElement;
      Method _getClassNameStackWalker;

      try {
        // StackWalker.getInstance(StackWalker.Option option)
        getStackWalker = Class.forName("java.lang.StackWalker").getMethod("getInstance",Class.forName("java.lang.StackWalker$Option"));
        // This invokes the previous method
        Set<Object> options = new HashSet<>();
        options.add(Class.forName("java.lang.StackWalker$Option").getEnumConstants()[0]);
        _stackWalker = getStackWalker.invoke(null, Class.forName("java.lang.StackWalker$Option").getEnumConstants()[0]);
        // This code makes the method StackWalker.walk
        _walk = _stackWalker.getClass().getMethod("walk", Function.class);
        // the method StackWalker.StackFrame.toStackTraceElement() instance method
        _toStackTraceElement = Class.forName("java.lang.StackWalker$StackFrame").getMethod("toStackTraceElement");
        // StackWalker.StackFrame.getClassName() intance method
        _getClassNameStackWalker = Class.forName("java.lang.StackWalker$StackFrame").getMethod("getClassName");
      }
      catch (Exception e) {
        LOG.log(Level.CONFIG, "Current configuration does not support fast implementation of getCaller(..)", e);
        _stackWalker = null;
        _walk = null;
        _toStackTraceElement = null;
        _getClassNameStackWalker = null;
      }

      stackWalker = _stackWalker;
      walk = _walk;
      toStackTraceElement = _toStackTraceElement;
      getClassNameStackWalker = _getClassNameStackWalker;

      // Assign JDK 8 fields to null
      javaLangAccess        =  null;
      getStackTraceDepth    =  null;
      getStackTraceElement  =  null;
    }
  }

  /** Gets the class+method that called the specified class.
   *  @param self The class that wants to know who called it.
   *  @return The caller or null if it could not be obtained.
   */
  public static StackTraceElement getCaller (Class<?> self) {
    return getCaller(self, null);
  }

  /**
   * This method will return the StackTraceElement to be logged based
   * on the first instance of the given Class in the StackTrace.
   * @param clazz The clazz that will be logged
   * @return The first instance of the StackTraceElement
   */
  public static StackTraceElement getCallerByClass (Class<?> clazz) {
    try {
      if (!CallStackUtilJDKSwap.JDK8_CALLSTACKUTIL) return CallStackUtilJDKSwap.getStackTraceElementJava9ByClass(clazz);
      else return getStackTraceElementJava8ByClass(clazz);
    } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
      LOG.log(Level.WARNING, "Error inspecting stack", e);
      return null;
    }
  }

  private static StackTraceElement getCallerJava8Code(Class<?> self, Class<?> expectType) {
    Throwable         t         = new Throwable();
    boolean           foundSelf = false;
    StackTraceElement save      = null;

    // ==== Fast Java 8 Version ====================================================================
    // This implementation is loosely based on code from OpenJDK and other online postings.
    // It is much faster than the other approach, but there is no guarantee of availability.
    // In general it works well with Java 8 (and earlier), but not Java 9 (and later).

      try {
        int depth = (Integer)getStackTraceDepth.invoke(javaLangAccess, t);

        for (int i = 0; i < depth; i++) {
          StackTraceElement elem  = (StackTraceElement)getStackTraceElement.invoke(javaLangAccess, t, i);
          String            clazz = elem.getClassName();

          try {
            if (clazz.equals(self.getName())) {
              foundSelf = true;
              continue;
            }
            if (!foundSelf) {
              continue; // keep moving down until self is found
            }
            if (clazz.startsWith("java.lang.reflect.") || clazz.startsWith("sun.reflect.")) {
              continue; // ignore reflection methods immediately below self
            }

            Class<?>           c      = Class.forName(clazz);
            IgnoreStackElement ignore = c.getAnnotation(IgnoreStackElement.class);

            if (isIgnored(ignore, elem)) {
              continue; // skip this method
            }

            if (expectType == null) {
              return elem;
            }

            if (save == null) {
              // save at the point where we would have otherwise returned
              save = elem;
            }

            if (clazz.startsWith("org.junit.") ||
                clazz.startsWith("org.gradle.api.internal.")) {
              // we have clearly gone too far
              return save;
            }
            if (expectType.isAssignableFrom(Class.forName(clazz))) {
              // found the one we are looking for
              return elem;
            }
          }
          catch (ClassNotFoundException e) {
            LOG.log(Level.WARNING, "Error inspecting stack", e);
            return (save == null)? elem : save; // Could not work further down the stack
          }
        }
        return save; // not found
      }
      catch (Exception e) {
        LOG.log(Level.WARNING, "Could not obtain caller using fast method", e);
        return null;
      }
  }
  /** Gets the class+method that called the specified class.
   *  @param self       The class that wants to know who called it.
   *  @param expectType The expected type (can be a superclass). Used when traversing the stack
   *                    looking for a stack trace element from a particular class type (null if
   *                    not applicable). For example, in the case of a {@code JAXBValidationException}
   *                    the {@code expectType} is {@code XMLObject} since we expect {@code XMLObject}
   *                    instances to throw a {@code JAXBValidationException} and are unconcerned
   *                    if the actual caller is a utility function called by the {@code XMLObject}
   *                    instance.
   *  @return The caller or null if it could not be obtained.
   */
  public static StackTraceElement getCaller (Class<?> self, Class<?> expectType) {


    if (javaLangAccess != null) {
      return getCallerJava8Code(self, expectType);
    }
    // If stackwalker != null, we can use that to do the logging with jdk9+
    if (stackWalker != null) {
      try {
        if (!CallStackUtilJDKSwap.JDK8_CALLSTACKUTIL) return CallStackUtilJDKSwap.getStackTraceElementJava9(self, expectType);
        else return getStackTraceElementJava8(self, expectType);
      }
      catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
        LOG.log(Level.WARNING, "Error inspecting stack", e);
        return null;
      }
    }

    // We should never get past this point

    // ==== Slow-but-Reliable Version ==============================================================
    // This implementation should always work. Note that this implementation that requires
    // creation of a Throwable is significantly faster than using Thread.getStackTrace()
    // since the Thread.getStackTrace() can be called by another thread and therefore
    // requires a hold-point while the stack trace is built; a Throwable doesn't need to
    // do this since it can only be created on the owning thread.
    { // <-- keep for consistency with above
      Throwable         t         = new Throwable();
      boolean           foundSelf = false;
      StackTraceElement save      = null;
      { // <-- keep for consistency with above
        for (StackTraceElement elem : t.getStackTrace()) {
          try {
            String clazz = elem.getClassName();

            if (clazz.equals(self.getName())) {
              foundSelf = true;
              continue;
            }
            if (!foundSelf) {
              continue; // keep moving down until self is found
            }
            if (clazz.startsWith("java.lang.reflect.") || clazz.startsWith("sun.reflect.")) {
              continue; // ignore reflection methods immediately below self
            }

            Class<?>           c      = Class.forName(clazz);
            IgnoreStackElement ignore = c.getAnnotation(IgnoreStackElement.class);

            if (isIgnored(ignore, elem)) {
              continue; // skip this method
            }

            if (expectType == null) {
              return elem;
            }
            if (save == null) {
              // save at the point where we would have otherwise returned
              save = elem;
            }

            if (clazz.startsWith("org.junit.") ||
                clazz.startsWith("org.gradle.api.internal.")) {
              // we have clearly gone too far
              return save;
            }
            if (expectType.isAssignableFrom(Class.forName(clazz))) {
              // found the one we are looking for
              return elem;
            }
          }
          catch (ClassNotFoundException e) {
            LOG.log(Level.WARNING, "Error inspecting stack", e);
            return (save == null)? elem : save; // Could not work further down the stack
          }
        }
        return save; // not found
      }
    }
  }

  /** Indicates if a given {@link StackTraceElement} points to code that is in a method that
   *  is marked with {@link IgnoreStackElement}.
   *  @param ignore The applicable {@link IgnoreStackElement} (can be null).
   *  @param elem   The applicable {@link StackTraceElement} (not null).
   *  @return true if this {@link StackTraceElement} should be ignored, false otherwise.
   */
  static boolean isIgnored (IgnoreStackElement ignore, StackTraceElement elem) {
//    System.out.println(ignore);

    if (ignore == null) return false;
    String[] methods = ignore.methods();

    if ((methods.length == 1) && (methods[0].equals("*"))) {
      return true;
    }

    String methodName = elem.getMethodName();
    for (String name : methods) {
      if (name.equals(methodName)) return true;
    }
    return false;
  }

  /** Indicates that certain methods (and/or the entire class) should be ignored by
   *  {@link #getCaller(Class,Class)}. This is generally used to flag methods that
   *  have special meaning for the purposes of logging.
   */
  @Documented
  @Inherited
  @Retention(value=RetentionPolicy.RUNTIME)
  public static @interface IgnoreStackElement {
    /** Ignore the specified methods in this class.
     *  @return The names of the methods to ignore or a single-element array with "*" as the value
     *          to ignore all methods in the class.
     */
    public String[] methods () default { "*" };
  }

  private static boolean matchHelper(Object f, Class<?> self, IsFoundHelper ifh) {
    if (self == null) return true;
    try {
      String clazz;

      clazz = (String) getClassNameStackWalker.invoke(f);

      if (!ifh.pastSelf && clazz.equals(self.getName())) {
        ifh.isFound = true; // This will be LoggingUtilities.class, the only one to skip
        return false;
      }
      if (!ifh.isFound) return false;
      ifh.pastSelf = true;
      ifh.elem = (StackTraceElement) toStackTraceElement.invoke(f);
      return true;
    } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
      LOG.log(Level.CONFIG, "Current configuration does not support fast implementation of getCaller(..)", e);
    }
    return true;
  }

  private static boolean matchExpectType(Object f, Class<?> expectType, IsFoundHelper ifh) {
    if (expectType == null) return true;
    try {
      if (expectType.isAssignableFrom(Class.forName((String)getClassNameStackWalker.invoke(f)))) {
        // found the one we are looking for
        ifh.elem = (StackTraceElement) toStackTraceElement.invoke(f);
        return true;
      }
    } catch (ClassNotFoundException | IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }
    return false;
  }

  /**
   *
   * @param f Stackwalker.Stackframe
   * @return true if it
   */
  private static boolean shouldIgnore (Object f) {
    try {
      String className;
      boolean dwInstanceOfStackTraceElement = f instanceof StackTraceElement;
      if (dwInstanceOfStackTraceElement) className= ((StackTraceElement)f).getClassName();
      else className = (String) getClassNameStackWalker.invoke(f);
      Class<?> clazz = Class.forName(className);
      IgnoreStackElement ignore = clazz.getAnnotation(IgnoreStackElement.class);
      if (isIgnored(ignore, dwInstanceOfStackTraceElement?(StackTraceElement) f: (StackTraceElement) toStackTraceElement.invoke(f))) {
        return true; // ignore this one
      } else {
        return false;
      }
    } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException | ClassNotFoundException e) {
      e.printStackTrace();
    }
    return false;
  }

  private static StackTraceElement getStackTraceElementJava8 (Class<?> self, Class<?> expectType)
      throws IllegalAccessException, IllegalArgumentException, InvocationTargetException {
    IsFoundHelper ifh = new IsFoundHelper();
    walk.invoke(stackWalker,
        (Function<Stream<Object>, Object>) s -> s
            .filter(f->dropWhileHelperReflection(f, self, expectType, ifh)) // drops until we meet the class that called
            .findFirst()
            .orElse(null));
    return ifh.elem;
  }

  private static boolean dropWhileHelperReflection(Object f, Class<?> self, Class<?> expectType, IsFoundHelper ifh) {
    String clazz;
    try {
      clazz = (String) getClassNameStackWalker.invoke(f);
      if (!ifh.pastSelf && clazz.equalsIgnoreCase(self.getName())) {
        ifh.isFound = true;
        return false;
      }
      if (!ifh.isFound) return false;
      ifh.pastSelf=true; // If we get here, we are done looking skipping the calling class, stop comparing strings
      Class<?> c = Class.forName(clazz);
      IgnoreStackElement ignore = c.getAnnotation(IgnoreStackElement.class);
      StackTraceElement elem = (StackTraceElement) toStackTraceElement.invoke(f);
      if (isIgnored(ignore, elem)) {
        return false; // skip this method
      }
      if (expectType == null) {
        ifh.elem = elem;
        return true;
      }
      if (ifh.save == null) {
        // save at the point where we would have otherwise returned
        ifh.save = elem;
      }

      if (clazz.startsWith("org.junit.") ||
          clazz.startsWith("org.gradle.api.internal.")
          || clazz.startsWith("junit.framework.")) {
        // we have clearly gone too far
        ifh.elem = ifh.save;
        return true;
      }
      if (expectType.isAssignableFrom(Class.forName(clazz))) {
        // found the one we are looking for
        ifh.elem = elem;
        return true;
      }
      //
    } catch (Exception e) {
      LOG.log(Level.CONFIG, "Current configuration does not support fast implementation of getCaller(..)", e);
    }
    return false;
  }

  private static StackTraceElement getStackTraceElementJava8ByClass (Class<?> self) throws IllegalAccessException, IllegalArgumentException, InvocationTargetException {
    Object elem = null;
    if(stackWalker!=null) { // This is java 9 code
      elem = walk.invoke(stackWalker, (Function<Stream<Object>, Object>) s -> s
          .filter(f -> {
            try {
              return getClassNameStackWalker.invoke(f).equals(self.getName());
            } catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
              // TODO Auto-generated catch block
              e.printStackTrace();
            }
            return false;
          })
          .findFirst()
          .orElse(null));
      return elem == null ? null : (StackTraceElement) toStackTraceElement.invoke(elem);
    }
    else {
      return getCaller(self,  null);
    }
  }
}
