JavaMessenger.java

/*
 * SPDX-FileCopyrightText: 2025 kaumei.io
 * SPDX-License-Identifier: Apache-2.0
 */
package io.kaumei.jdbc.anno.ctx;

import io.kaumei.jdbc.anno.ProcessorException;
import io.kaumei.jdbc.anno.msg.Msg;
import io.kaumei.jdbc.annotation.config.JdbcLogLevel;
import org.jspecify.annotations.Nullable;

import javax.annotation.processing.Messager;
import javax.annotation.processing.ProcessingEnvironment;
import javax.lang.model.element.*;
import javax.tools.Diagnostic;
import java.util.Arrays;
import java.util.Collection;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;

import static io.kaumei.jdbc.anno.Processor.OPTION_KEY_LOG_LEVEL;
import static java.util.Objects.requireNonNull;

public class JavaMessenger {

    // ----- service
    private final Context ctx;
    private final Messager messager;

    // ----- state
    private JdbcLogLevel.LogLevel logState = JdbcLogLevel.LogLevel.ERROR;

    public JavaMessenger(Context ctx, ProcessingEnvironment env) {
        this.ctx = requireNonNull(ctx);
        this.messager = requireNonNull(env.getMessager());
        this.logState = getLogLevel(env.getOptions());
    }

    // ------------------------------------------------------------------------

    void updateLogState(JdbcLogLevel.LogLevel level) {
        this.logState = level;
    }

    public boolean isDebugEnabled() {
        return this.logState == JdbcLogLevel.LogLevel.DEBUG;
    }

    public <T extends Element> void acceptWithDebugFlag(T elem, Consumer<T> consumer) {
        var oldLoggerState = enableDebug(elem);
        try {
            consumer.accept(elem);
        } finally {
            this.logState = oldLoggerState;
        }
    }

    public <T extends Element, R> R applyWithDebugFlag(T elem, Function<T, R> func) {
        var oldLoggerState = enableDebug(elem);
        try {
            return func.apply(elem);
        } finally {
            this.logState = oldLoggerState;
        }
    }

    private JdbcLogLevel.LogLevel enableDebug(Element context) {
        var element = context;
        while (element != null && element.getKind() != ElementKind.PACKAGE) {
            if (this.ctx.JDBC_DEBUG.hasAnno(element)) {
                var old = logState;
                logState = JdbcLogLevel.LogLevel.DEBUG;
                return old;
            }
            element = element.getEnclosingElement();
        }
        return logState;
    }

    // ------------------------------------------------------------------------

    public void always(CharSequence msg, Object... args) {
        this.messager.printMessage(Diagnostic.Kind.NOTE, format(msg, args));
    }

    // ------------------------------------------------------------------------

    /**
     * @param msg the message
     * @param args key value pairs
     */
    public void error(CharSequence msg, Object... args) {
        this.messager.printMessage(Diagnostic.Kind.ERROR, format(msg, args));
    }

    public void error(@Nullable Element e, CharSequence msg, Object... args) {
        this.messager.printMessage(Diagnostic.Kind.ERROR, format(msg, args), e);
    }

    public void error(Exception exp) {
        var expStr = formatException(exp);
        if (exp instanceof ProcessorException pe) {
            this.messager.printMessage(Diagnostic.Kind.ERROR, expStr, pe.element());
        } else {
            this.messager.printMessage(Diagnostic.Kind.ERROR, expStr);
        }
    }

    // ------------------------------------------------------------------------

    public void warn(CharSequence msg, Object... args) {
        if (JdbcLogLevel.LogLevel.WARN.isEnabled(logState)) {
            this.messager.printMessage(Diagnostic.Kind.WARNING, format(msg, args));
        }
    }

    public void warn(@Nullable Element e, CharSequence msg, Object... args) {
        if (JdbcLogLevel.LogLevel.WARN.isEnabled(logState)) {
            this.messager.printMessage(Diagnostic.Kind.WARNING, format(msg, args), e);
        }
    }

    // ------------------------------------------------------------------------

    public void info(CharSequence msg, Object... args) {
        if (JdbcLogLevel.LogLevel.INFO.isEnabled(logState)) {
            this.messager.printMessage(Diagnostic.Kind.NOTE, format(msg, args));
        }
    }

    public void info(@Nullable Element e, CharSequence msg, Object... args) {
        if (JdbcLogLevel.LogLevel.INFO.isEnabled(logState)) {
            this.messager.printMessage(Diagnostic.Kind.NOTE, format(msg, args), e);
        }
    }

    // ------------------------------------------------------------------------

    public void debug(CharSequence msg, Object... args) {
        if (JdbcLogLevel.LogLevel.DEBUG.isEnabled(logState)) {
            this.messager.printMessage(Diagnostic.Kind.NOTE, format(msg, args));
        }
    }

    public void debug(@Nullable Element e, CharSequence msg, Object... args) {
        if (JdbcLogLevel.LogLevel.DEBUG.isEnabled(logState)) {
            this.messager.printMessage(Diagnostic.Kind.NOTE, format(msg, args), e);
        }
    }

    // ------------------------------------------------------------------------

    public CharSequence format(CharSequence msg, @Nullable Object... args) {
        if (args.length == 0) {
            return msg;
        }
        var kv = new StringBuilder(msg);
        kv.append(' ');

        var sep = false;
        int i = 0;
        while (i + 1 < args.length) {
            addKeyValue(kv, sep, args[i], args[i + 1]);
            i += 2;
            sep = true;
        }

        if (i < args.length) {
            if (args[i] instanceof Exception e) {
                kv.append(": ");
                kv.append(e);
                addStackTrace(kv, e.getStackTrace());
            } else if (args[i] != null) {
                addKeyValue(kv, sep, null, args[i]);
            }
        }
        return kv;
    }

    // ------------------------------------------------------------------------

    public JdbcLogLevel.LogLevel getLogLevel(Map<String, String> map) {
        var value = map.get(OPTION_KEY_LOG_LEVEL);
        if (value != null) {
            try {
                return JdbcLogLevel.LogLevel.valueOf(value.toUpperCase());
            } catch (Exception e) {
                this.warn("Found invalid log level. Use default ERROR.", "value", value);
            }
        }
        return JdbcLogLevel.LogLevel.ERROR;
    }

    private void addKeyValue(StringBuilder sb, boolean sep, @Nullable Object key, @Nullable Object value) {
        if (sep) {
            sb.append(", ");
        }
        if (key != null) {
            sb.append(key);
            sb.append('=');
        }
        addValueTo(sb, value);
    }

    private void addValueTo(StringBuilder sb, @Nullable Object value) {
        if (value instanceof Msg.Messages messages) {
            sb.append(messages);
        } else if (value instanceof Iterable<?> iterable) {
            boolean sep = false;
            sb.append('[');
            if (iterable instanceof Collection<?> c) {
                sb.append("size:");
                sb.append(c.size());
                sep = true;
            }
            for (Object item : iterable) {
                if (sep) {
                    sb.append(", ");
                } else {
                    sep = true;
                }
                if (item instanceof Element e) {
                    sb.append(e.getSimpleName());
                } else {
                    sb.append(item);
                }
            }
            sb.append(']');
        } else if (value instanceof Element e) {
            this.addQualifiedName(sb, e);
        } else {
            sb.append(value);
        }
    }

    private void addQualifiedName(StringBuilder sb, Element element) {
        if (element instanceof PackageElement p) {
            sb.append(p.getQualifiedName());
        } else if (element instanceof TypeElement t) {
            sb.append(t.getQualifiedName());
        } else if (element instanceof ExecutableElement e) {
            sb.append(e.getEnclosingElement());
            sb.append(".");
            sb.append(e.getSimpleName());
        } else {
            sb.append(element.getSimpleName());
        }
        var pos = this.ctx.sourcePosition(element);
        if (pos != null) {
            sb.append(' ').append(pos);
        }
    }

    public static String formatException(Exception exception) {
        var sb = new StringBuilder();
        sb.append(exception);
        addStackTrace(sb, exception.getStackTrace());
        return sb.toString();
    }

    public static String printStackTrace() {
        var sb = new StringBuilder();
        var stackTrace = Thread.currentThread().getStackTrace();
        addStackTrace(sb, Arrays.copyOfRange(stackTrace, 2, stackTrace.length));
        return sb.toString();
    }

    private static void addStackTrace(StringBuilder sb, StackTraceElement[] list) {
        int lastIndex = list.length - 1;
        while (lastIndex > 0 && !list[lastIndex - 1].getClassName().startsWith("io.kaumei")) {
            lastIndex--;
        }
        for (int i = 0; i <= lastIndex; i++) {
            sb.append("\n\tat ");
            sb.append(list[i]);
        }
        sb.append('\n');
    }

}