TargetMethod.java

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

import com.palantir.javapoet.CodeBlock;
import com.palantir.javapoet.MethodSpec;
import io.kaumei.jdbc.CodeGenerationException;
import io.kaumei.jdbc.JdbcEmptyResultSetException;
import io.kaumei.jdbc.JdbcUnexpectedRowException;
import io.kaumei.jdbc.anno.ProcessorException;
import io.kaumei.jdbc.anno.ctx.ConfigService;
import io.kaumei.jdbc.anno.ctx.Context;
import io.kaumei.jdbc.anno.jdbc2java.ColumnIndex;
import io.kaumei.jdbc.anno.jdbc2java.Jdbc2JavaConverter;
import io.kaumei.jdbc.anno.model.OptionalFlag;
import io.kaumei.jdbc.anno.model.SQLNameDV;
import io.kaumei.jdbc.anno.msg.JdbcMsg;
import io.kaumei.jdbc.anno.msg.Msg;
import io.kaumei.jdbc.annotation.config.JdbcNoMoreRows;
import io.kaumei.jdbc.annotation.config.JdbcNoRows;
import org.jspecify.annotations.Nullable;

import javax.lang.model.element.ExecutableElement;
import javax.lang.model.type.TypeMirror;
import java.sql.Types;
import java.util.*;

import static java.util.Objects.requireNonNull;

public final class TargetMethod {

    private static final Set<String> RESERVED = Set.of("con", "index", "plCount", "result", "resultSet", "row", "rs", "sql", "sqlSb", "stmt", "supplier");

    private final Context ctx;
    private final @Nullable ExecutableElement method;
    // ----- state
    private final CodeBlock.Builder code = CodeBlock.builder();
    private final Set<String> usedNames = new HashSet<>();
    private final Map<String, String> parameterNames = new HashMap<>();

    public TargetMethod(Context ctx, ExecutableElement method) {
        this.ctx = ctx;
        this.method = method;
        this.usedNames.addAll(RESERVED);

        var conflicts = new ArrayList<String>();
        // keep all non-conflicting names
        for (var param : method.getParameters()) {
            var name = param.getSimpleName().toString();
            if (this.usedNames.add(name)) {
                parameterNames.put(name, name);
            } else {
                conflicts.add(name);
            }
        }
        for (var param : conflicts) {
            parameterNames.put(param, tempVarName(param + "Param"));
        }
    }

    TargetMethod(TargetMethod parent) {
        this.ctx = parent.ctx;
        this.method = null;
        this.usedNames.addAll(parent.usedNames);
        this.parameterNames.putAll(parent.parameterNames);
    }

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

    @Override
    public String toString() {
        return this.code.build().toString();
    }

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

    CodeBlock build() {
        return code.build();
    }

    MethodSpec build(Msg.Messages messages, String header) {
        var methodBuilder = ctx.kaumeiJdbcGenerator.createMethodBuilder(requireNonNull(method), this);
        if (messages.hasMessages()) {
            var errorCode = CodeBlock.builder();
            var msg = render(messages, header);
            errorCode.addStatement("throw new $T($S)", CodeGenerationException.class, msg);
            methodBuilder.addCode(errorCode.build());
        } else {
            methodBuilder.addCode(code.build());
        }
        return methodBuilder.build();
    }

    private String render(Msg.Messages messages, String header) {
        var sb = new StringBuilder();
        sb.append(header).append("\n");

        boolean reasons = true;
        for (var message : messages) {
            if (message instanceof JdbcMsg.InvalidParamConverter
                    || message instanceof JdbcMsg.InvalidReturnConverter) {
                continue;
            } else if (reasons) {
                sb.append("\nReason(s):");
                reasons = false;
            }
            sb.append("\n* ").append(message.text());
        }

        var params = true;
        for (var message : messages) {
            if (message instanceof JdbcMsg.InvalidParamConverter) {
                if (params) {
                    sb.append("\nJava2Jdbc parameter converter(s):");
                    params = false;
                }
                sb.append('\n');
                sb.append(message.text());
            }
        }

        var returnType = true;
        for (var message : messages) {
            if (message instanceof JdbcMsg.InvalidReturnConverter) {
                if (returnType) {
                    sb.append("\nJdbc2Java return converter:");
                    returnType = false;
                }
                sb.append('\n');
                sb.append(message.text());

            }
        }
        return sb.toString();
    }

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

    public void addComment(String message, Object... args) {
        var msg = this.ctx.logger.format(message, args);
        this.ctx.logger.debug("Add comment to code:", msg);
        this.code.add("// " + msg + "\n");
    }

    public void indent() {
        this.code.indent();
    }

    public void unindent() {
        this.code.unindent();
    }

    public void beginControlFlow(String controlFlow, Object... args) {
        this.code.beginControlFlow(controlFlow, args);
    }

    public void nextControlFlow(String controlFlow, Object... args) {
        this.code.nextControlFlow(controlFlow, args);
    }

    public void endControlFlow() {
        this.code.endControlFlow();
    }

    public void addStatement(String format, Object... args) {
        this.code.addStatement(format, args);
    }

    public void add(String format, Object... args) {
        this.code.add(format, args);
    }

    public void addCodeBlock(CodeBlock codeBlock) {
        this.code.add(codeBlock);
    }

    public void addStatement(CodeBlock codeBlock) {
        this.code.addStatement(codeBlock);
    }

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

    public String paramName(String sourceName) {
        return requireNonNull(parameterNames.get(sourceName));
    }

    public CodeBlock paramCodeBlock(String sourceName) {
        return CodeBlock.of("$N", paramName(sourceName));
    }

    void addIfAnnotationIsPresent(String format, ConfigService.@Nullable OptionValue<?> result) {
        if (result instanceof ConfigService.OptionValue.Dynamic<?> d) {
            this.code.add(checkValue(d));
            this.code.addStatement(format, accessValue(d));
        } else if (result instanceof ConfigService.OptionValue.Static<?> s) {
            this.code.addStatement(format, accessValue(s));
        }
    }

    public CodeBlock checkValue(ConfigService.OptionValue.Dynamic<?> result) {
        var varName = paramName(result.paramName());
        var option = result.option();
        var value = option.defaultValue();
        var checkBlock = option.minValue() != null
                ? CodeBlock.of("$N < $L", varName, option.minValue())
                : value.getClass().isEnum()
                  ? CodeBlock.of("$N == $T.$L", varName, value.getClass(), value)
                  : CodeBlock.of("$N == $L", varName, value);
        return CodeBlock.builder()
                .beginControlFlow("if($L)", checkBlock)
                .addStatement("throw new $T($S + $L)", IllegalArgumentException.class, "Invalid value for " + varName + ": ", varName)
                .endControlFlow().build();
    }

    public CodeBlock accessValue(ConfigService.OptionValue<?> result) {
        if (result instanceof ConfigService.OptionValue.Dynamic<?> d) {
            return paramCodeBlock(d.paramName());
        } else if (result instanceof ConfigService.OptionValue.Static<?> s) {
            var value = s.value();
            return value.getClass().isEnum()
                    ? CodeBlock.of("$T.$L", value.getClass(), value)
                    : CodeBlock.of("$L", value);
        }
        throw new IllegalArgumentException("Unknown: " + result);  // sanity-check
    }

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

    public CodeBlock setNull(TypeMirror type, CodeBlock columnIndex) {
        // @formatter:off
        return switch (type.getKind()) { // will never cover all branches in black box test: JaCoCo:ignore
            case BOOLEAN -> CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "BOOLEAN");
            case BYTE ->    CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "TINYINT");
            case SHORT ->   CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "SMALLINT");
            case INT ->     CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "INTEGER");
            case LONG ->    CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "BIGINT");
            case CHAR ->    CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "CHAR");
            case FLOAT ->   CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "REAL");
            case DOUBLE ->  CodeBlock.of("stmt.setNull($L, $T.$L)", columnIndex, Types.class, "DOUBLE");
            default -> throw new ProcessorException("type must be primitive"); // sanity-check
        };
        // @formatter:on
    }

    CodeBlock lambda(SQLNameDV jdbcName, OptionalFlag optional, Jdbc2JavaConverter converter) {
        var lambda = new TargetMethod(this);
        if (converter.isColumn()) {
            ColumnIndex index;
            if (!jdbcName.hasName()) {
                index = ColumnIndex.ofValue(1);
            } else {
                index = ColumnIndex.ofVariable("index", jdbcName);
                this.addStatement("var $N = resultSet.findColumn($S)", index.columnIndexVar(), index.columnName());
            }
            lambda.add("(rs) -> {\n");
            lambda.indent();
            converter.addColumnByIndex(lambda, "row", index, optional);
            lambda.addStatement("return row");
            lambda.unindent();
            lambda.add("}");
        } else {
            lambda.add("(rs) -> {\n");
            lambda.indent();
            converter.addResultSetToRow(lambda, "row", optional);
            lambda.addStatement("return row");
            lambda.unindent();
            lambda.add("}");
        }
        return lambda.build();
    }

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

    public void addThrowColumnWasNull(ColumnIndex index) {
        if (index.hasColumnName()) {
            this.code.addStatement("throw new $T($S + $L)", NullPointerException.class, "JDBC column was null on name '" + index.columnName() + "' with index: ", index.columnIndexVar());
        } else {
            this.code.addStatement("throw new $T($S + $L)", NullPointerException.class, "JDBC column was null on index: ", index.columnIndexVar());
        }
    }

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

    void addCheckNoRows(JdbcNoRows.Kind noRows, OptionalFlag optional) {
        switch (noRows) { // JaCoCo:ignore
            case THROW_EXCEPTION -> {
                this.beginControlFlow("if(!rs.next())");
                this.addStatement("throw new $T()", JdbcEmptyResultSetException.class);
                this.endControlFlow();
            }
            case RETURN_NULL -> {
                this.beginControlFlow("if(!rs.next())");
                if (optional.isOptionalType()) {
                    // JdbcRow is the only case for this
                    this.addStatement("return $T.empty()", Optional.class);
                } else {
                    this.addStatement("return null");
                }
                this.endControlFlow();
            }
            default -> throw new ProcessorException("invalid converter: " + noRows); // sanity-check
        }
    }

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

    void addCheckNoMoreRows(JdbcNoMoreRows.Kind noMoreRows) {
        switch (noMoreRows) { // JaCoCo:ignore
            case THROW_EXCEPTION -> {
                this.beginControlFlow("if(rs.next())");
                this.addStatement("throw new $T()", JdbcUnexpectedRowException.class);
                this.endControlFlow();
            }
            case IGNORE -> this.addStatement("// ignore more results ");
            default ->
                    throw new ProcessorException("invalid converter: " + noMoreRows); // sanity-check
        }
    }

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

    public String tempVarName(String varName) {
        if (usedNames.add(varName)) {
            return varName;
        }
        for (int i = 2; ; i++) {
            var candidate = varName + "_" + i;
            if (usedNames.add(candidate)) {
                return candidate;
            }
        }
    }

    public CodeBlock tempVarName(CodeBlock varName) {
        var varName0 = varName.toString().replaceAll("[^\\p{IsAlphabetic}\\d_$]+", "_");
        var name = tempVarName(varName0);
        return CodeBlock.of(name);
    }

}