SourceMethodParameter.java

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

import io.kaumei.jdbc.anno.ctx.ConfigService;
import io.kaumei.jdbc.anno.ctx.Context;
import io.kaumei.jdbc.anno.msg.JdbcMsg;
import io.kaumei.jdbc.anno.msg.Msg;
import io.kaumei.jdbc.anno.store.SearchKey;
import io.kaumei.jdbc.anno.utils.SqlDV;
import io.kaumei.jdbc.anno.utils.SqlTokenizer;
import org.jspecify.annotations.Nullable;

import javax.lang.model.element.*;
import javax.lang.model.type.DeclaredType;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

import static java.util.Objects.requireNonNull;

public class SourceMethodParameter {

    static SourceMethodParameter of(Context ctx, Msg.Builder messages, ExecutableElement method) {
        return of0(ctx, messages, method, null);
    }

    static SourceMethodParameter of(Context ctx, Msg.Builder messages, ExecutableElement method, SqlDV sql) {
        return of0(ctx, messages, method, sql);
    }

    private static SourceMethodParameter of0(Context ctx, Msg.Builder messages, ExecutableElement method, @Nullable SqlDV sql) {
        var name2param = name2param(ctx, method);
        if (sql != null) {
            var tokens = SqlTokenizer.parse(sql);
            var hasCollections = new ParameterContext(ctx, new Java2JdbcResolver(ctx),
                    messages, name2param, tokens)
                    .processAndUpdate();
            return new SourceMethodParameter(name2param, hasCollections, tokens, sql);
        }
        return new SourceMethodParameter(name2param, null, null, null);
    }

    private static Map<String, Param> name2param(Context ctx, ExecutableElement method) {
        var name2param = new HashMap<String, Param>();
        for (var paramElem : method.getParameters()) {
            var name = paramElem.getSimpleName().toString();
            var map = AnnoMap.of(ctx, paramElem);
            map.consume(ctx.JDBC_CONVERTER_NAME); // special case: this anno is consuming in the converter
            var param = new ParamUnknown(map, paramElem);
            name2param.put(name, param);
        }
        return name2param;
    }

    private record ParameterContext(Context ctx,
                                    Java2JdbcResolver resolver,
                                    Msg.Builder messages,
                                    Map<String, Param> name2param,
                                    List<SqlTokenizer.Token> tokens) {
        boolean processAndUpdate() {
            var hasCollections = new AtomicBoolean(false);
            int plCount = 0;
            for (var token : tokens) {
                if (token instanceof SqlTokenizer.UnnamedParameter ignore) {
                    messages.add(JdbcMsg.sqlPlaceholderUnnamedUnsupported());
                } else if (token instanceof SqlTokenizer.Parameter tokenParam) {
                    if (!tokenParam.path().isEmpty()) {
                        messages.add(JdbcMsg.sqlPlaceholderPathUnsupported(tokenParam.name()));
                        continue;
                    }
                    var name = tokenParam.root();
                    if (!name2param.containsKey(name)) {
                        messages.add(JdbcMsg.sqlPlaceholderNotFound(name));
                        continue;
                    }
                    var param0 = name2param.get(name);
                    if (param0 instanceof ParamUnknown param) {
                        if (token instanceof SqlTokenizer.SingleParameter) {
                            singleParameter(name, param);
                        } else if (token instanceof SqlTokenizer.AllValues) {
                            valuesParameter(name, param, hasCollections);
                        } else if (token instanceof SqlTokenizer.AllNames) {
                            namesParameter(name, param);
                        }
                    } else if (param0 instanceof ParamSingle && token instanceof SqlTokenizer.SingleParameter) {
                        // ok
                    } else if (param0 instanceof ParamValues &&
                            (token instanceof SqlTokenizer.AllValues || token instanceof SqlTokenizer.AllNames)) {
                        // ok
                    } else {
                        var pName = shapeName(param0);
                        var tName = shapeName(token);
                        messages.add(JdbcMsg.sqlPlaceholderUsedWithDifferentShapes(name, pName, tName));
                        continue;
                    }

                    // update plCount
                    var resolved = name2param.get(name);
                    if (resolved instanceof ParamSingle || resolved instanceof ParamArray || resolved instanceof ParamList) {
                        plCount++;
                    } else if (resolved instanceof ParamComposite pc && token instanceof SqlTokenizer.AllValues) {
                        plCount += pc.names.length;
                    }
                    if (plCount > ctx.kaumeiConfig.maxTotalPlaceholders()) {
                        messages.add(JdbcMsg.sqlPlaceholderCountExceedsLimit(plCount, ctx.kaumeiConfig.maxTotalPlaceholders()));
                    }
                }
            } // for
            return hasCollections.get();
        }

        String shapeName(SqlTokenizer.Token token) {
            if (token instanceof SqlTokenizer.SingleParameter) {
                return "single";
            } else if (token instanceof SqlTokenizer.AllValues) {
                return "values";
            } else if (token instanceof SqlTokenizer.AllNames) {
                return "names";
            } else {
                return "unknown";
            }
        }

        String shapeName(Param token) {
            if (token instanceof ParamSingle) {
                return "single";
            } else if (token instanceof ParamValues) {
                return "values";
            } else {
                return "unknown";
            }
        }

        void singleParameter(String name, ParamUnknown param) {
            var result = resolver.single(param.element);
            if (result.hasMessages()) {
                messages.add(JdbcMsg.invalidParam(param.element(), result.messages()));
            } else {
                var value = result.value();
                name2param.put(name, param.asSingle(value.searchRequest(), value.optional()));
            }
        }

        void valuesParameter(String name, ParamUnknown param, AtomicBoolean hasCollections) {
            var converterName = ctx.JDBC_CONVERTER_NAME.getAnnoOpt(param.element);
            if (converterName != null && converterName.isBlankName()) {
                messages.add(JdbcMsg.invalidParam(param.element(),
                        JdbcMsg.annotationValueMustNotBeBlank("@JdbcConverterName")));
                return;
            } else if (converterName != null) {
                messages.add(JdbcMsg.javaMethodParameterJdbcConverterNameUnsupportedForSqlPlaceholder(
                        name, ":" + name + ".{values}", "values"));
                return;
            }
            var result = resolver.values(param.element);
            if (result.hasMessages()) {
                messages.add(JdbcMsg.invalidParam(param.element(), result.messages()));
            } else if (result.value().kind() == Java2JdbcResolver.Kind.ARRAY) {
                var value = result.value();
                name2param.put(name, param.asArray(value.searchRequest(), value.optional(), value.cmpOptional()));
                hasCollections.set(true);
            } else if (result.value().kind() == Java2JdbcResolver.Kind.LIST) {
                var value = result.value();
                name2param.put(name, param.asList(value.searchRequest(), value.optional(), value.cmpOptional()));
                hasCollections.set(true);
            } else if (result.value().kind() == Java2JdbcResolver.Kind.RECORD) {
                valuesOfRecord(name, param);
            } else {
                throw new RuntimeException("!!!");
            }
        }

        void valuesOfRecord(String name, ParamUnknown param) {
            if (!(param.element.asType() instanceof DeclaredType declared)
                    || !(declared.asElement() instanceof TypeElement recordElement)
                    || recordElement.getKind() != ElementKind.RECORD) {
                throw new RuntimeException("unreachable");
            }
            var recordComponents = recordElement.getRecordComponents();
            var names = new String[recordComponents.size()];
            var searchRequests = new SearchKey[recordComponents.size()];
            var vars = new RecordComponentElement[recordComponents.size()];
            var optionals = new OptionalFlag[recordComponents.size()];
            for (int i = 0; i < recordComponents.size(); i++) {
                var cmp = recordComponents.get(i);
                var jdbcName = ctx.JDBC_NAME.getAnnoOpt(cmp);
                if (jdbcName != null && jdbcName.hasMessages()) {
                    messages.add(JdbcMsg.invalidParam(param.element(), jdbcName.messages()));
                    return;
                }
                var converterName = ctx.JDBC_CONVERTER_NAME.getAnnoOpt(cmp);
                if (converterName != null && converterName.isBlankName()) {
                    messages.add(JdbcMsg.invalidParam(param.element(),
                            JdbcMsg.annotationValueMustNotBeBlank("@JdbcConverterName")));
                    return;
                }
                vars[i] = cmp;
                names[i] = cmp.getSimpleName().toString();
                searchRequests[i] = SearchKey.of(ctx, cmp);
                optionals[i] = ctx.optionalFlag(cmp);
            }
            name2param.put(name, param.asComposite(names, searchRequests, vars, optionals));
        }

        void namesParameter(String name, ParamUnknown param) {
            var converterName = ctx.JDBC_CONVERTER_NAME.getAnnoOpt(param.element);
            if (converterName != null && converterName.isBlankName()) {
                messages.add(JdbcMsg.invalidParam(param.element(),
                        JdbcMsg.annotationValueMustNotBeBlank("@JdbcConverterName")));
                return;
            } else if (converterName != null) {
                messages.add(JdbcMsg.javaMethodParameterJdbcConverterNameUnsupportedForSqlPlaceholder(
                        name, ":" + name + ".{names}", "names"));
                return;
            }
            var result = resolver.values(param.element);
            if (result.hasMessages()) {
                messages.add(JdbcMsg.invalidParam(param.element(), result.messages()));
            } else if (result.value().kind() == Java2JdbcResolver.Kind.RECORD) {
                valuesOfRecord(name, param);
            } else {
                messages.add(JdbcMsg.invalidParam(param.element(), JdbcMsg.JAVA_TO_JDBC_RESOLVER_NAMES_UNSUPPORTED_TYPES));
            }
        }
    }

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

    private final Map<String, Param> name2param;
    private final @Nullable Boolean hasCollectionsOpt;
    private final @Nullable List<SqlTokenizer.Token> tokensOpt;
    private final @Nullable SqlDV sqlOpt;

    private SourceMethodParameter(Map<String, Param> name2param,
                                  @Nullable Boolean hasCollectionsOpt,
                                  @Nullable List<SqlTokenizer.Token> tokensOpt,
                                  @Nullable SqlDV sqlOpt) {
        this.name2param = name2param;
        this.hasCollectionsOpt = hasCollectionsOpt;
        this.tokensOpt = tokensOpt;
        this.sqlOpt = sqlOpt;
    }

    public boolean hasCollections() {
        return requireNonNull(hasCollectionsOpt);
    }

    public List<SqlTokenizer.Token> tokens() {
        return requireNonNull(tokensOpt);
    }

    public SqlDV sql() {
        return requireNonNull(sqlOpt);
    }

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

    public void addUnusedAnno(Msg.Builder messages) {
        for (var entry : name2param.entrySet()) {
            var paramAnno = entry.getValue().annoMap();
            if (paramAnno.hasUnused()) {
                var name = entry.getKey();
                messages.add(JdbcMsg.javaMethodParameterHasUnusedAnnotations(name, paramAnno.unused()));
            }
        }
    }

    @Nullable
    public <T> String paramNameFor(ConfigService.Option<T> option) {
        for (var entry : name2param.entrySet()) {
            if (entry.getValue().annoMap().getIfDefault(option) != null) {
                return entry.getKey();
            }
        }
        return null;
    }

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

    // for test purpose only
    Set<String> parameterNames() {
        return name2param.keySet();
    }

    public void forEach(Consumer<Param> consumer) {
        for (var param : this.name2param.values()) {
            consumer.accept(param);
        }
    }

    public void forEachSearchRequest(Consumer<SearchKey> consumer) {
        for (var p : this.name2param.values()) {
            for (var sr : p.allSearchRequests()) {
                consumer.accept(sr);
            }
        }
    }

    public Param get(String name) {
        return requireNonNull(name2param.get(name));
    }

    public ParamSingle getSingle(String name) {
        var param = name2param.get(name);
        if (param instanceof SourceMethodParameter.ParamSingle ps) {
            return ps;
        }
        throw new RuntimeException("Parameter '" + name + "' is not a composite: " + param);
    }

    public ParamValues getValues(String name) {
        var param = name2param.get(name);
        if (param instanceof SourceMethodParameter.ParamValues pv) {
            return pv;
        }
        throw new RuntimeException("Parameter '" + name + "' is not a composite: " + param);
    }

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

    public sealed interface Param permits ParamSingle, ParamUnknown, ParamValues {
        AnnoMap annoMap();

        VariableElement element();

        default Set<SearchKey> allSearchRequests() {
            return Collections.emptySet();
        }
    }

    public record ParamUnknown(AnnoMap annoMap, VariableElement element) implements Param {
        ParamArray asArray(SearchKey searchRequest, OptionalFlag optional, OptionalFlag cmpOptional) {
            return new ParamArray(annoMap, element, searchRequest, optional, cmpOptional);
        }

        ParamList asList(SearchKey searchRequest, OptionalFlag optional, OptionalFlag cmpOptional) {
            return new ParamList(annoMap, element, searchRequest, optional, cmpOptional);
        }

        ParamSingle asSingle(SearchKey searchRequest, OptionalFlag optional) {
            return new ParamSingle(annoMap, element, searchRequest, optional);
        }

        ParamComposite asComposite(String[] names, SearchKey[] searchRequests, RecordComponentElement[] variable, OptionalFlag[] optionals) {
            return new ParamComposite(annoMap, element, names, searchRequests, variable, optionals);
        }
    }

    public record ParamSingle(AnnoMap annoMap, VariableElement element,
                              SearchKey searchRequest,
                              OptionalFlag optional) implements Param {
        @Override
        public Set<SearchKey> allSearchRequests() {
            return Set.of(searchRequest);
        }

    }

    public sealed interface ParamValues extends Param permits ParamArray, ParamList, ParamComposite {
    }

    public record ParamArray(AnnoMap annoMap, VariableElement element, SearchKey searchRequest,
                             OptionalFlag optional,
                             OptionalFlag cmpOptional) implements ParamValues {
        @Override
        public Set<SearchKey> allSearchRequests() {
            return Set.of(searchRequest);
        }
    }

    public record ParamList(AnnoMap annoMap, VariableElement element,
                            SearchKey searchRequest,
                            OptionalFlag optional,
                            OptionalFlag cmpOptional) implements ParamValues {
        @Override
        public Set<SearchKey> allSearchRequests() {
            return Set.of(searchRequest);
        }
    }

    public record ParamComposite(AnnoMap annoMap, VariableElement element, String[] names,
                                 SearchKey[] searchRequests,
                                 RecordComponentElement[] variable,
                                 OptionalFlag[] optionals) implements ParamValues {
        @Override
        public Set<SearchKey> allSearchRequests() {
            return Set.of(searchRequests);
        }
    }
}