gson-compile/gson-compile-processor/src/main/java/io/gitlab/jfronny/gson/compile/processor/adapter/impl/MapAdapter.java

package io.gitlab.jfronny.gson.compile.processor.adapter.impl;

import com.squareup.javapoet.*;
import io.gitlab.jfronny.gson.compile.processor.Cl;
import io.gitlab.jfronny.gson.compile.processor.core.TypeHelper;
import io.gitlab.jfronny.gson.compile.processor.adapter.Adapter;

import javax.lang.model.element.ElementKind;
import javax.lang.model.element.Modifier;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeMirror;
import java.io.IOException;
import java.util.*;

public class MapAdapter extends Adapter<MapAdapter.Hydrated> {
    @Override
    public Hydrated instantiate() {
        return new Hydrated();
    }

    public class Hydrated extends Adapter<MapAdapter.Hydrated>.Hydrated {
        private static final List<Class<? extends Map>> SUPPORTED = List.of(LinkedHashMap.class, HashMap.class, TreeMap.class);

        private DeclaredType type;
        private TypeName implType;
        private TypeMirror componentType1;
        private TypeMirror componentType2;

        @Override
        public boolean applies() {
            return type != null;
        }

        @Override
        protected void afterHydrate() {
            type = TypeHelper.asDeclaredType(super.type);
            componentType1 = componentType2 = null;
            if (type == null) return;
            List<? extends TypeMirror> typeArguments = type.getTypeArguments();
            if (typeArguments.size() != 2) {
                type = null;
            } else {
                componentType1 = typeArguments.get(0);
                if (!isValidKey(componentType1)) {
                    type = null;
                    componentType1 = null;
                    return;
                }
                componentType2 = typeArguments.get(1);
                String ts = TypeHelper.asDeclaredType(typeUtils.erasure(type)).asElement().toString();
                if (Map.class.getCanonicalName().equals(ts)) {
                    implType = TypeName.get(SUPPORTED.get(0));
                    return;
                }
                for (Class<?> klazz : SUPPORTED) {
                    if (klazz.getCanonicalName().equals(ts)) {
                        implType = TypeName.get(klazz);
                        return;
                    }
                }
                type = null;
                componentType1 = componentType2 = null;
            }
        }

        private boolean isValidKey(TypeMirror tm) {
            if (tm.toString().equals(String.class.getCanonicalName())) return true;
            if (tm.toString().equals(UUID.class.getCanonicalName())) return true;
            if (unbox(tm).getKind().isPrimitive()) return true;
            if (isEnum(tm)) return true;
            return false;
        }

        private void generateConvertKey(CodeBlock.Builder kode) {
            if (componentType1.toString().equals(String.class.getCanonicalName())) {
                kode.add("name");
            }
            else if (componentType1.toString().equals(UUID.class.getCanonicalName())) {
                kode.add("name == null ? null : $T.fromString(name)", UUID.class);
            }
            else if (unbox(componentType1).getKind().isPrimitive()) {
                kode.add("name == null ? null : " + switch (unbox(componentType1).getKind()) {
                    case BOOLEAN -> "Boolean.parseBoolean(name)";
                    case BYTE -> "Byte.parseByte(name)";
                    case SHORT -> "Short.parseShort(name)";
                    case INT -> "Integer.parseInt(name)";
                    case LONG -> "Long.parseLong(name)";
                    case CHAR -> "name.length() == 0 ? '\\0' : name.charAt(0)";
                    case FLOAT -> "Float.parseFloat(name)";
                    case DOUBLE -> "Double.parseDouble(name)";
                    default -> throw new IllegalArgumentException("Unsupported primitive: " + unbox(componentType1).getKind());
                });
            }
            else if (isEnum(componentType1)) {
                String methodName = "readName$" + name;
                boolean found = false;
                for (MethodSpec spec : klazz.methodSpecs) {
                    if (spec.name.equals(methodName)) {
                        found = true;
                        break;
                    }
                }
                if (!found) {
                    klazz.addMethod(
                            MethodSpec.methodBuilder(methodName)
                                    .addModifiers(Modifier.PRIVATE, Modifier.STATIC)
                                    .returns(TypeName.get(componentType1))
                                    .addParameter(String.class, "value")
                                    .addCode(
                                            CodeBlock.builder()
                                                    .beginControlFlow("for ($1T t : $1T.values())", componentType1)
                                                    .addStatement("if (t.name().equals(value)) return t")
                                                    .endControlFlow()
                                                    .addStatement("return null")
                                                    .build()
                                    )
                                    .build()
                    );
                }
                kode.add("$N(name)", methodName);
            }
        }

        private boolean isEnum(TypeMirror tm) {
            DeclaredType declared = TypeHelper.asDeclaredType(tm);
            return declared != null && declared.asElement().getKind() == ElementKind.ENUM;
        }

        @Override
        public void generateWrite(Runnable writeGet) {
            code.addStatement("writer.beginObject()");
            code.add("for ($T.Entry<$T, $T> $N : (", Map.class, componentType1, componentType2, argName);
            writeGet.run();
            code.beginControlFlow(").entrySet())")
                    .beginControlFlow("if ($N.getKey() != null || writer.getSerializeNulls())", argName)
                    .add("writer.name(");
            if (isEnum(componentType1)) {
                code.add("$N.getKey() == null ? \"null\" : $N.getKey().name()", argName, argName);
            } else {
                code.add("$T.toString($N.getKey())", Objects.class, argName);
            }
            code.addStatement(")")
                    .addStatement("$T value$N = $N.getValue()", componentType2, argName, argName)
                    .beginControlFlow("if (value$N == null)", argName)
                    .addStatement("if (writer.getSerializeNulls()) writer.nullValue()")
                    .nextControlFlow("else");
            generateWrite(code, componentType2, "value" + argName, componentType2.getAnnotationMirrors(), () -> code.add("value" + argName));
            code.endControlFlow().endControlFlow().endControlFlow().addStatement("writer.endObject()");
        }

        @Override
        public void generateRead() {
            CodeBlock.Builder kode = CodeBlock.builder();
            kode.addStatement("$T map = new $T<>()", typeName, implType);

            kode.addStatement("reader.beginObject()")
                    .beginControlFlow("while (reader.hasNext())")
                    .addStatement("String name = reader.nextName()")
                    .beginControlFlow("if (reader.peek() == $T.NULL)", Cl.GSON_TOKEN)
                    .addStatement("reader.nextNull()")
                    .add("map.put(");
            generateConvertKey(kode);
            kode.addStatement(", null)")
                    .nextControlFlow("else")
                    .add("map.put(");
            generateConvertKey(kode);
            kode.add(", ");
            generateRead(kode, componentType2, argName, componentType2.getAnnotationMirrors());
            kode.addStatement(")")
                    .endControlFlow()
                    .endControlFlow()
                    .addStatement("reader.endObject()")
                    .addStatement("return map");

            String methodName = "read$" + name;
            klazz.addMethod(
                    MethodSpec.methodBuilder(methodName)
                            .addModifiers(Modifier.PRIVATE, Modifier.STATIC)
                            .returns(typeName)
                            .addParameter(Cl.GSON_READER, "reader")
                            .addException(IOException.class)
                            .addCode(kode.build())
                            .build()
            );
            code.add("$N(reader)", methodName);
        }
    }
}