package org.jfleet.avro;

import org.apache.avro.Schema;
import org.apache.avro.SchemaBuilder;
import org.apache.avro.SchemaBuilder.FieldAssembler;
import org.apache.avro.SchemaBuilder.FieldTypeBuilder;
import org.jfleet.ColumnInfo;
import org.jfleet.EntityInfo;

public class AvroSchemaBuilder {
    private final EntityInfo entityInfo;

    public AvroSchemaBuilder(EntityInfo entityInfo) {
        this.entityInfo = entityInfo;
    }

    public Schema build() {
        FieldAssembler<Schema> fields = SchemaBuilder.record(entityInfo.getEntityClass().getName())
                .namespace(entityInfo.getEntityClass().getPackage().getName()).fields();

        for (ColumnInfo columnInfo : entityInfo.getColumns()) {
            fields = getFieldSchema(columnInfo, fields);
        }
        return fields.endRecord();
    }

    private FieldAssembler<Schema> getFieldSchema(ColumnInfo columnInfo, FieldAssembler<Schema> fields) {
        FieldTypeBuilder<Schema> beginType = fields.name(columnInfo.getColumnName()).type();
        boolean isPrimitive = columnInfo.getFieldType().isPrimitive();

        switch (columnInfo.getFieldType().getFieldType()) {
            case STRING:
            case ENUMSTRING:
                return buildAsString(beginType);
            case ENUMORDINAL:
                return buildAsInteger(beginType, false);
            case INT:
            case SHORT:
            case BYTE:
                return buildAsInteger(beginType, isPrimitive);
            case DOUBLE:
                return buildAsDouble(beginType, isPrimitive);
            case LONG:
                return buildAsLong(beginType, isPrimitive);
            case FLOAT:
                return buildAsFloat(beginType, isPrimitive);
            case BOOLEAN:
                return buildAsBoolean(beginType, isPrimitive);
            default:
                throw new UnsupportedTypeException(String.format("Unsupported type: %s", columnInfo.getFieldType().getFieldType()));
        }

    }

    private FieldAssembler<Schema> buildAsString(FieldTypeBuilder<Schema> type) {
        return type.unionOf().stringType().and().nullType().endUnion().noDefault();
    }

    private FieldAssembler<Schema> buildAsBoolean(FieldTypeBuilder<Schema> type, boolean isPrimitive) {
        if (isPrimitive) {
            return type.booleanType().noDefault();
        }
        return type.unionOf().booleanType().and().nullType().endUnion().noDefault();
    }

    private FieldAssembler<Schema> buildAsFloat(FieldTypeBuilder<Schema> type, boolean isPrimitive) {
        if (isPrimitive) {
            return type.floatType().noDefault();
        }
        return type.unionOf().floatType().and().nullType().endUnion().noDefault();
    }

    private FieldAssembler<Schema> buildAsLong(FieldTypeBuilder<Schema> type, boolean isPrimitive) {
        if (isPrimitive) {
            return type.longType().noDefault();
        }
        return type.unionOf().longType().and().nullType().endUnion().noDefault();
    }

    private FieldAssembler<Schema> buildAsInteger(FieldTypeBuilder<Schema> type, boolean isPrimitive) {
        if (isPrimitive) {
            return type.intType().noDefault();
        }
        return type.unionOf().intType().and().nullType().endUnion().noDefault();
    }

    private FieldAssembler<Schema> buildAsDouble(FieldTypeBuilder<Schema> type, boolean isPrimitive) {
        if (isPrimitive) {
            return type.doubleType().noDefault();
        }
        return type.unionOf().doubleType().and().nullType().endUnion().noDefault();
    }
}

@@ -19,16 +19,52 @@
 .avro;%0A%0A
+import java.util.function.Function;%0A
 import o
@@ -116,32 +116,91 @@
 .SchemaBuilder;%0A
+import org.apache.avro.SchemaBuilder.BaseFieldTypeBuilder;%0A
 import org.apach
@@ -232,24 +232,75 @@
 dAssembler;%0A
+import org.apache.avro.SchemaBuilder.FieldDefault;%0A
 import org.a
@@ -717,20 +717,16 @@
         
-    
 .namespa
@@ -1377,30 +1377,51 @@
 rn b
-uildAsString(beginType
+eginType.nullable().stringType().noDefault(
 );%0A 
@@ -1468,38 +1468,30 @@
      return 
-buildAsInteger
+ofType
 (beginType, 
@@ -1495,16 +1495,47 @@
 e, false
+, BaseFieldTypeBuilder::intType
 );%0A     
@@ -1621,30 +1621,22 @@
  return 
-buildAsInteger
+ofType
 (beginTy
@@ -1642,32 +1642,63 @@
 ype, isPrimitive
+, BaseFieldTypeBuilder::intType
 );%0A            c
@@ -1732,28 +1732,21 @@
  return 
-buildAsDoubl
+ofTyp
 e(beginT
@@ -1753,32 +1753,66 @@
 ype, isPrimitive
+, BaseFieldTypeBuilder::doubleType
 );%0A            c
@@ -1844,27 +1844,22 @@
  return 
-buildAsLong
+ofType
 (beginTy
@@ -1865,32 +1865,64 @@
 ype, isPrimitive
+, BaseFieldTypeBuilder::longType
 );%0A            c
@@ -1955,28 +1955,22 @@
  return 
-buildAsFloat
+ofType
 (beginTy
@@ -1976,32 +1976,65 @@
 ype, isPrimitive
+, BaseFieldTypeBuilder::floatType
 );%0A            c
@@ -2069,30 +2069,22 @@
  return 
-buildAsBoolean
+ofType
 (beginTy
@@ -2098,16 +2098,51 @@
 rimitive
+, BaseFieldTypeBuilder::booleanType
 );%0A     
@@ -2208,16 +2208,37 @@
 ception(
+%0A                    
 String.f
@@ -2368,1518 +2368,277 @@
 ma%3E 
-buildAsString(FieldTypeBuilder%3CSchema%3E type) %7B%0A        return type.unionOf().stringType().and().nullType().endUnion().noDefault();%0A    %7D%0A%0A    private FieldAssembler%3CSchema%3E buildAsBoolean(FieldTypeBuilder%3CSchema%3E type, boolean isPrimitive) %7B%0A        if (isPrimitive) %7B%0A            return type.booleanType().noDefault();%0A        %7D%0A        return type.unionOf().booleanType().and().nullType().endUnion().noDefault();%0A    %7D%0A%0A    private FieldAssembler%3CSchema%3E buildAsFloat(FieldTypeBuilder%3CSchema%3E type, boolean isPrimitive) %7B%0A        if (isPrimitive) %7B%0A            return type.floatType().noDefault();%0A        %7D%0A        return type.unionOf().floatType().and().nullType().endUnion().noDefault();%0A    %7D%0A%0A    private FieldAssembler%3CSchema%3E buildAsLong(FieldTypeBuilder%3CSchema%3E type, boolean isPrimitive) %7B%0A        if (isPrimitive) %7B%0A            return type.longType().noDefault();%0A        %7D%0A        return type.unionOf().longType().and().nullType().endUnion().noDefault();%0A    %7D%0A%0A    private FieldAssembler%3CSchema%3E buildAsInteger(FieldTypeBuilder%3CSchema%3E type, boolean isPrimitive) %7B%0A        if (isPrimitive) %7B%0A            return type.intType().noDefault();%0A        %7D%0A        return type.unionOf().intType().and().nullType().endUnion().noDefault();%0A    %7D%0A%0A    private FieldAssembler%3CSchema%3E buildAsDouble(FieldTypeBuilder%3CSchema%3E type, boolean isPrimitive) %7B%0A        if (isPrimitive) %7B%0A            return type.doubleType().noDefault();%0A        %7D%0A        return type.unionOf().doubleType().and().nullType().endUnion(
+ofType(FieldTypeBuilder%3CSchema%3E type, boolean isPrimitive,%0A        Function%3CBaseFieldTypeBuilder%3CSchema%3E, FieldDefault%3CSchema, ?%3E%3E typeMapper) %7B%0A        BaseFieldTypeBuilder%3CSchema%3E forType = isPrimitive ? type : type.nullable();%0A        return typeMapper.apply(forType
 ).no
@@ -2654,10 +2654,11 @@
 ;%0A    %7D%0A
+%0A
 %7D%0A

Simplify build schema code using builder