GenerateJdbcBatchUpdate.java

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

import com.palantir.javapoet.ClassName;
import com.palantir.javapoet.MethodSpec;
import com.palantir.javapoet.TypeSpec;
import io.kaumei.jdbc.JdbcException;
import io.kaumei.jdbc.anno.ctx.ConfigService;
import io.kaumei.jdbc.anno.ctx.Context;
import io.kaumei.jdbc.anno.model.JdbcTypeKind;
import io.kaumei.jdbc.anno.model.SourceMethod;
import io.kaumei.jdbc.anno.model.SourceMethodParameter;
import io.kaumei.jdbc.anno.msg.JdbcMsg;
import io.kaumei.jdbc.anno.msg.Msg;
import io.kaumei.jdbc.impl.JdbcBatchImpl;
import org.jspecify.annotations.Nullable;

import javax.lang.model.element.*;
import javax.lang.model.type.TypeKind;
import java.sql.PreparedStatement;
import java.sql.SQLException;

import static io.kaumei.jdbc.anno.utils.AnnoUtils.requireState;

public class GenerateJdbcBatchUpdate implements GenerateJdbc {
    // ----- services
    private final Context ctx;
    // ----- state
    private final SourceMethod sourceMethod;
    private final TargetMethod targetMethod;
    private final Msg.Builder messages;
    private final KaumeiClassBuilder parent;

    GenerateJdbcBatchUpdate(Context ctx, KaumeiClassBuilder parent, SourceMethod sourceMethod) {
        this.ctx = ctx;
        this.sourceMethod = sourceMethod;
        this.targetMethod = new TargetMethod(this.ctx, sourceMethod.method());
        this.messages = Msg.builder();
        this.parent = parent;
    }

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

    public MethodSpec generateMethod() {
        this.ctx.logger.debug("---- @JdbcBatchUpdate ----");

        var returnType = this.sourceMethod.returnType();
        this.ctx.logger.debug("returnType", returnType);
        var optReason = returnType.optional().checkNonNullOrUnspecified();
        if (optReason != null) {
            this.messages.add(JdbcMsg.jdbcBatchRequiresNonNullOrUnspecifiedReturnType(returnType.optional()));
        } else if (returnType.kind() != JdbcTypeKind.KAUMEI_JDBC_BATCH) {
            this.messages.add(JdbcMsg.jdbcBatchRequiresBatchInterfaceReturnType(returnType));
        }

        var batchSize = requireState(this.sourceMethod.jdbcBatchSize(), "batchSize must not be null");
        var batchType = this.ctx.asElementOpt(returnType.type());
        var updateMethod = batchType == null ? null : getUpdateMethod(batchType);
        if (updateMethod == null || this.messages.hasMessages()) {
            return this.targetMethod.build(this.messages.build(), "@JdbcBatchUpdate method invalid");
        }

        updateMethod.parameter().forEach(param -> {
            if (param instanceof SourceMethodParameter.ParamArray || param instanceof SourceMethodParameter.ParamList) {
                messages.add(JdbcMsg.jdbcBatchUpdateMethodMustNotUseDynamicCollectionExpansion(param.element()));
            }
        });
        if (messages.hasMessages()) {
            return this.targetMethod.build(this.messages.build(), "@JdbcBatchUpdate method invalid");
        }

        // ------------
        var batchClassName = updateMethod.parent().getSimpleName() + ctx.kaumeiConfig.generatedClassSuffix();
        var batchMethod = this.parent.containsClass(batchClassName);
        if (batchMethod == null) {
            var batchClass = TypeSpec.classBuilder(batchClassName)
                    .addModifiers(Modifier.STATIC)
                    .superclass(ClassName.get(JdbcBatchImpl.class))
                    .addSuperinterface(ClassName.get(updateMethod.parent()))
                    .addMethod(MethodSpec.constructorBuilder()
                            .addParameter(PreparedStatement.class, "stmt")
                            .addParameter(int.class, "batchSize")
                            .addStatement("super(stmt, batchSize)")
                            .build());
            batchMethod = new TargetMethod(this.ctx, updateMethod.method());

            if (updateMethod.method().getReturnType().getKind() != TypeKind.VOID) {
                this.messages.add(JdbcMsg.jdbcBatchUpdateMethodMustReturnVoidMessage(
                        updateMethod.parent().getSimpleName(),
                        updateMethod.method().getSimpleName()));
                return this.targetMethod.build(this.messages.build(), "@JdbcUpdate method invalid");
            }

            batchMethod.beginControlFlow("try");
            this.ctx.kaumeiJdbcGenerator.processParameter(messages, updateMethod, batchMethod);

            batchMethod.addStatement("super.addBatch()");

            batchMethod.nextControlFlow("catch ($T e)", SQLException.class);
            batchMethod.addStatement("throw new $T(e.getMessage(), e)", JdbcException.class);
            batchMethod.endControlFlow();

            batchClass.addMethod(batchMethod.build(Msg.empty(), "@JdbcUpdate method invalid"));
            this.parent.addClass(batchClassName, batchClass.build(), batchMethod);
        }

        if (batchSize instanceof ConfigService.OptionValue.Dynamic<?> d) {
            var batchSizeBlock = targetMethod.accessValue(d);
            targetMethod.beginControlFlow("if($L < 1)", batchSizeBlock);
            targetMethod.addStatement("throw new $T($S + $L)", IllegalArgumentException.class, "Invalid value for " + d.paramName() + ": ", batchSizeBlock);
            targetMethod.endControlFlow();
        }

        targetMethod.beginControlFlow("try");
        targetMethod.addStatement("var con = supplier.getConnection()");
        // ----
        targetMethod.addCodeBlock(this.ctx.kaumeiJdbcGenerator.buildSqlVariable(updateMethod, batchMethod, "sql"));
        targetMethod.addStatement("var stmt = con.prepareStatement(sql)");
        targetMethod.addIfAnnotationIsPresent("stmt.setQueryTimeout($L)", this.sourceMethod.jdbcQueryTimeout());

        var batchSizeBlock = targetMethod.accessValue(batchSize);
        targetMethod.addStatement("return new $N(stmt, $L)", batchClassName, batchSizeBlock);
        // ----
        targetMethod.nextControlFlow("catch ($T e)", SQLException.class);
        targetMethod.addStatement("throw new $T(e.getMessage(), e)", JdbcException.class);
        targetMethod.endControlFlow();

        this.messages.add(sourceMethod.unusedAnno());
        var unusedUpdateAnnotations = updateMethod.unusedAnno();
        if (unusedUpdateAnnotations.hasMessages()) {
            this.messages.add(unusedUpdateAnnotations);
        }
        return this.targetMethod.build(this.messages.build(), "@JdbcBatchUpdate method invalid");
    }

    @Nullable SourceMethod getUpdateMethod(@Nullable Element batchType) {
        if (!(batchType instanceof TypeElement)) {
            this.messages.add(JdbcMsg.jdbcBatchTypeMustBeInterfaceMessage(String.valueOf(batchType)));
            return null;
        }
        var batchTypeElement = (TypeElement) batchType;
        if (!batchTypeElement.getKind().isInterface()) {
            this.messages.add(JdbcMsg.jdbcBatchTypeMustBeInterfaceMessage(batchTypeElement.getQualifiedName()));
            return null;
        }
        if (!batchTypeElement.getEnclosingElement().equals(this.parent.type())) {
            this.messages.add(JdbcMsg.jdbcBatchTypeMustBeNestedInFactoryInterface(batchTypeElement.getQualifiedName()));
            return null;
        }
        var methods = batchTypeElement.getEnclosedElements();
        if (methods.size() == 1) {
            var method0 = methods.get(0);
            if (!method0.getModifiers().contains(Modifier.STATIC)
                    && !method0.getModifiers().contains(Modifier.DEFAULT)
                    && method0.getKind() == ElementKind.METHOD
                    && method0 instanceof ExecutableElement updateMethod) {
                var result = ctx.sourceMethodService.getOpt(updateMethod);
                if (result != null && result.entryPoint().updateOpt() != null) {
                    if (result.messages().hasMessages()) {
                        this.messages.add(result.messages());
                    }
                    return result;
                }
            }
        }
        this.messages.add(JdbcMsg.jdbcBatchTypeMustDeclareExactlyOneUpdateMethod(batchTypeElement.getQualifiedName()));
        return null;
    }

}