package com.github.gfx.static_gson;

import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.FieldSpec;
import com.squareup.javapoet.ParameterizedTypeName;
import com.squareup.javapoet.TypeName;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import javax.lang.model.element.Modifier;

public class TypeRegistry {

    private final Map<TypeName, FieldSpec> registry = new HashMap<>();

    public TypeRegistry() {
    }

    private static FieldSpec createField(TypeName type) {
        return FieldSpec.builder(Types.getTypeAdapter(type), toName(type))
                .addModifiers(Modifier.PRIVATE, Modifier.FINAL)
                .build();
    }

    private static CodeBlock toInitializer(TypeName type) {
        return CodeBlock.of("gson.getAdapter($T.get($L))", Types.TypeToken, toTypeExpr(type));
    }

    private static CodeBlock toTypeExpr(TypeName type) {
        if (type instanceof ParameterizedTypeName) {
            ParameterizedTypeName parameterizedType = (ParameterizedTypeName) type;
            CodeBlock.Builder typeArgsExpr = CodeBlock.builder();
            for (int i = 0; i < parameterizedType.typeArguments.size(); i++) {
                TypeName t = parameterizedType.typeArguments.get(i);
                if (i != 0) {
                    typeArgsExpr.add(", $L", toTypeExpr(t));
                } else {
                    typeArgsExpr.add("$L", toTypeExpr(t));
                }
            }
            return CodeBlock.of("$T.newParameterizedTypeWithOwner(null, $T.class, $L)",
                            Types.$Gson$Types, parameterizedType.rawType, typeArgsExpr.build());
        } else {
            return CodeBlock.of("$T.class", type);
        }
    }

    private static String toName(TypeName type) {
        return "$" + type.toString().replaceAll("[^a-zA-Z0-9?]+", "\\$");
    }

    public FieldSpec getField(TypeName type) {
        FieldSpec field = registry.get(type);
        if (field == null) {
            field = createField(type);
            registry.put(type, field);
        }
        return field;
    }

    public Collection<FieldSpec> getFields() {
        return registry.values();
    }

    public CodeBlock getFieldInitialization() {
        CodeBlock.Builder block = CodeBlock.builder();
        registry.forEach((type, field) -> {
            if (type instanceof ParameterizedTypeName) {
                block.addStatement("$N = ($T) $L",
                        field, Types.getTypeAdapter(type), toInitializer(type));
            } else {
                block.addStatement("$N = $L",
                        field, toInitializer(type));
            }
        });
        return block.build();
    }
}

@@ -463,238 +463,8 @@
  %7D%0A%0A
-    private static FieldSpec createField(TypeName type) %7B%0A        return FieldSpec.builder(Types.getTypeAdapter(type), toName(type))%0A                .addModifiers(Modifier.PRIVATE, Modifier.FINAL)%0A                .build();%0A    %7D%0A%0A
     
@@ -1696,46 +1696,69 @@
     
-FieldSpec field = registry.get(type);%0A
+return registry.compute(type, (typeName, fieldSpec) -%3E %7B%0A    
     
@@ -1770,16 +1770,20 @@
 f (field
+Spec
  == null
@@ -1802,71 +1802,190 @@
     
-field = createField(type);%0A            registry.put(type, fie
+    return FieldSpec.builder(Types.getTypeAdapter(typeName), toName(typeName))%0A                        .addModifiers(Modifier.PRIVATE, Modifier.FINAL)%0A                        .bui
 ld
+(
 );%0A 
@@ -1991,26 +1991,41 @@
         
+   
-%7D%0A
+ %7D else %7B%0A    
         
 return f
@@ -2012,24 +2012,28 @@
             
+    
 return field
@@ -2032,16 +2032,46 @@
 rn field
+Spec;%0A            %7D%0A        %7D)
 ;%0A    %7D%0A

use Map#compute()