package org.xbib.logging.util;

import java.lang.module.ModuleDescriptor;
import java.util.EnumSet;
import java.util.Iterator;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Stream;
import org.xbib.logging.ExtLogRecord;
import org.xbib.logging.LogContext;

public final class StackWalkerUtil {

    static final StackWalker WALKER = new GetStackWalkerAction().run();

    private StackWalkerUtil() {
    }

    public static Class<?> findCallingClass(Set<ClassLoader> rejectClassLoaders) {
        return WALKER.walk(new FindFirstWalkFunction(rejectClassLoaders));
    }

    public static LogContext logContextFinder(Set<ClassLoader> rejectClassLoaders,
                                              final Function<ClassLoader, LogContext> finder) {
        return WALKER.walk(new FindAllWalkFunction(rejectClassLoaders, finder));
    }

    public static void calculateCaller(ExtLogRecord logRecord) {
        WALKER.walk(new CallerCalcFunction(logRecord));
    }

    private static void calculateJdkModule(final ExtLogRecord logRecord, final Class<?> clazz) {
        final Module module = clazz.getModule();
        if (module != null) {
            logRecord.setSourceModuleName(module.getName());
            final ModuleDescriptor descriptor = module.getDescriptor();
            if (descriptor != null) {
                final Optional<ModuleDescriptor.Version> optional = descriptor.version();
                if (optional.isPresent()) {
                    logRecord.setSourceModuleVersion(optional.get().toString());
                } else {
                    logRecord.setSourceModuleVersion(null);
                }
            }
        }
    }

    private static final class CallerCalcFunction implements Function<Stream<StackWalker.StackFrame>, Void> {
        private final ExtLogRecord logRecord;

        CallerCalcFunction(final ExtLogRecord logRecord) {
            this.logRecord = logRecord;
        }

        @Override
        public Void apply(final Stream<StackWalker.StackFrame> stream) {
            final String loggerClassName = logRecord.getLoggerClassName();
            final Iterator<StackWalker.StackFrame> iterator = stream.iterator();
            boolean found = false;
            while (iterator.hasNext()) {
                final StackWalker.StackFrame frame = iterator.next();
                final Class<?> clazz = frame.getDeclaringClass();
                if (clazz.getName().equals(loggerClassName)) {
                    // next entry could be the one we want!
                    found = true;
                } else if (found) {
                    logRecord.setSourceClassName(frame.getClassName());
                    logRecord.setSourceMethodName(frame.getMethodName());
                    logRecord.setSourceFileName(frame.getFileName());
                    logRecord.setSourceLineNumber(frame.getLineNumber());
                    calculateJdkModule(logRecord, clazz);
                    return null;
                }
            }
            logRecord.setUnknownCaller();
            return null;
        }
    }

    private static final class GetStackWalkerAction {
        GetStackWalkerAction() {
        }

        public StackWalker run() {
            return StackWalker.getInstance(EnumSet.of(StackWalker.Option.RETAIN_CLASS_REFERENCE));
        }
    }

    private static final class FindFirstWalkFunction implements Function<Stream<StackWalker.StackFrame>, Class<?>> {
        private final Set<ClassLoader> rejectClassLoaders;

        FindFirstWalkFunction(final Set<ClassLoader> rejectClassLoaders) {
            this.rejectClassLoaders = rejectClassLoaders;
        }

        @Override
        public Class<?> apply(final Stream<StackWalker.StackFrame> stream) {
            final Iterator<StackWalker.StackFrame> iterator = stream.iterator();
            while (iterator.hasNext()) {
                final Class<?> clazz = iterator.next().getDeclaringClass();
                final ClassLoader classLoader = clazz.getClassLoader();
                if (!rejectClassLoaders.contains(classLoader)) {
                    return clazz;
                }
            }
            return null;
        }
    }

    private static final class FindAllWalkFunction implements Function<Stream<StackWalker.StackFrame>, LogContext> {
        private final Set<ClassLoader> rejectClassLoaders;
        private final Function<ClassLoader, LogContext> finder;

        FindAllWalkFunction(final Set<ClassLoader> rejectClassLoaders, final Function<ClassLoader, LogContext> finder) {
            this.rejectClassLoaders = rejectClassLoaders;
            this.finder = finder;
        }

        @Override
        public LogContext apply(final Stream<StackWalker.StackFrame> stream) {
            final Iterator<StackWalker.StackFrame> iterator = stream.iterator();
            while (iterator.hasNext()) {
                final Class<?> clazz = iterator.next().getDeclaringClass();
                final ClassLoader classLoader = clazz.getClassLoader();
                if (classLoader != null && !rejectClassLoaders.contains(classLoader)) {
                    final LogContext result = finder.apply(classLoader);
                    if (result != null) {
                        return result;
                    }
                }
            }
            return null;
        }
    }
}
