KaumeiClassBuilder.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.FieldSpec;
import com.palantir.javapoet.MethodSpec;
import com.palantir.javapoet.TypeSpec;
import io.kaumei.jdbc.anno.ProcessorException;
import io.kaumei.jdbc.anno.ctx.Context;
import io.kaumei.jdbc.core.CodeGenerationException;
import io.kaumei.jdbc.core.JdbcConnectionProvider;
import org.jspecify.annotations.Nullable;

import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.TypeMirror;
import java.util.HashMap;
import java.util.Map;

import static io.kaumei.jdbc.anno.Processor.GENERATED;
import static io.kaumei.jdbc.anno.utils.PrintStackTrace.appendStackTrace;

public class KaumeiClassBuilder {
    // ----- state
    private final TypeElement iface;
    private final String packageName;
    private final TypeSpec.Builder typeSpecBuilder;
    private final Map<String, TargetMethod> classNames = new HashMap<>();

    KaumeiClassBuilder(Context ctx, String packageName, TypeElement iface) {
        this.iface = iface;
        this.packageName = packageName;
        var simpleName = iface.getSimpleName() + ctx.kaumeiConfig.generatedClassSuffix();

        var generationAnno = ctx.JDBC_GENERATION.getAnnoOpt(iface);

        this.typeSpecBuilder = TypeSpec.classBuilder(simpleName)
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ClassName.get(iface));

        this.typeSpecBuilder.addAnnotation(GENERATED);
        // ----- add additional specified annotations
        if (generationAnno != null) {
            for (TypeMirror anno : generationAnno.classAnnotations()) {
                this.typeSpecBuilder.addAnnotation(ClassName.get(ctx.asTypeElement(anno)));
            }
        }

        // ----- Add supplier field
        typeSpecBuilder.addField(
                FieldSpec.builder(JdbcConnectionProvider.class, "supplier", Modifier.PRIVATE, Modifier.FINAL).build());


        var constructor = MethodSpec.constructorBuilder()
                .addModifiers(Modifier.PUBLIC)
                .addParameter(JdbcConnectionProvider.class, "supplier")
                .addStatement("this.supplier = $L", KaumeiLib.requireNonNull("supplier"));
        // ----- add additional specified annotations
        if (generationAnno != null) {
            for (var anno : generationAnno.constructorAnnotations()) {

                constructor.addAnnotation(ClassName.get(ctx.asTypeElement(anno)));
            }
        }

        // ----- Add constructor
        this.typeSpecBuilder.addMethod(constructor.build());
    }

    public TypeElement type() {
        return this.iface;
    }

    public TypeSpec build() {
        return this.typeSpecBuilder.build();
    }

    public String packageName() {
        return packageName;
    }

    @Nullable TargetMethod containsClass(String name) {
        return classNames.get(name);
    }

    void addClass(String name, TypeSpec typeSpec, TargetMethod targetMethod) {
        if (classNames.put(name, targetMethod) != null) { // sanity-check
            throw new ProcessorException("Name already known: " + name); // sanity-check
        }
        this.typeSpecBuilder.addType(typeSpec);
    }

    public void addMethod(MethodSpec methodSpec) {
        this.typeSpecBuilder.addMethod(methodSpec);
    }

    public void addMethodThrows(ExecutableElement method, Exception e) {
        var sb = new StringBuilder();
        sb.append("Annotation processing caught internal exception:\n");
        appendStackTrace(sb, e);
        this.addMethod(MethodSpec.overriding(method)
                .addStatement("throw new $T($S)", CodeGenerationException.class, sb.toString())
                .build());
        // sanity-check:off
    }

}