package io.quarkus.avro.runtime.graal;

import java.lang.reflect.Constructor;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.function.Function;

import org.apache.avro.Schema;
import org.apache.avro.generic.GenericData;
import org.apache.avro.io.DatumReader;
import org.apache.avro.io.ResolvingDecoder;
import org.apache.avro.util.WeakIdentityHashMap;

import com.oracle.svm.core.annotate.Alias;
import com.oracle.svm.core.annotate.RecomputeFieldValue;
import com.oracle.svm.core.annotate.Substitute;
import com.oracle.svm.core.annotate.TargetClass;

import io.quarkus.runtime.graal.JDK17OrLater;

@TargetClass(className = "org.apache.avro.generic.GenericDatumReader")
final class Target_org_apache_avro_generic_GenericDatumReader {

    @Alias
    private GenericData data;
    @Alias
    private Schema actual;
    @Alias
    private Schema expected;
    @Alias
    private DatumReader<?> fastDatumReader;
    @Alias
    private ResolvingDecoder creatorResolver;
    @Alias
    @RecomputeFieldValue(kind = RecomputeFieldValue.Kind.Reset)
    private Thread creator;

    @Alias
    @RecomputeFieldValue(kind = RecomputeFieldValue.Kind.FromAlias)
    private static ThreadLocal<Map<Schema, Map<Schema, ResolvingDecoder>>> RESOLVER_CACHE = ThreadLocal.withInitial(
            WeakIdentityHashMap::new);
    @Alias
    @RecomputeFieldValue(kind = RecomputeFieldValue.Kind.Reset)
    private Map<Schema, Class> stringClassCache;
    @Alias
    @RecomputeFieldValue(kind = RecomputeFieldValue.Kind.Reset)
    private Map<Class, Constructor> stringCtorCache;

    @Substitute
    protected Target_org_apache_avro_generic_GenericDatumReader(GenericData data) {
        this.fastDatumReader = null;
        this.creatorResolver = null;
        this.stringClassCache = new IdentityHashMap();
        this.stringCtorCache = new HashMap();
        this.data = data;
        this.creator = Thread.currentThread();
    }

}

@TargetClass(className = "org.apache.avro.reflect.ReflectionUtil", onlyWith = JDK17OrLater.class)
final class Target_org_apache_avro_reflect_ReflectionUtil {

    @Substitute
    public static <V, R> Function<V, R> getConstructorAsFunction(Class<V> parameterClass, Class<R> clazz) {
        // Cannot use the method handle approach as it uses ProtectionDomain which got removed.
        try {
            Constructor<R> constructor = clazz.getConstructor(parameterClass);
            return new Function<V, R>() {
                @Override
                public R apply(V v) {
                    try {
                        return constructor.newInstance(v);
                    } catch (Exception e) {
                        throw new IllegalStateException("Unable to create new instance for " + clazz, e);
                    }
                }
            };
        } catch (Throwable t) {
            // if something goes wrong, do not provide a Function instance
            return null;
        }
    }
}

class AvroSubstitutions {
}

@@ -75,68 +75,8 @@
 or;%0A
-import java.util.HashMap;%0Aimport java.util.IdentityHashMap;%0A
 impo
@@ -336,16 +336,77 @@
 HashMap;
+%0Aimport org.apache.avro.util.internal.ThreadLocalWithInitial;
 %0A%0Aimport
@@ -1301,10 +1301,9 @@
 ocal
-.w
+W
 ithI
@@ -1308,17 +1308,16 @@
 hInitial
-(
 %0A       
@@ -1321,16 +1321,20 @@
         
+.of(
 WeakIden
@@ -1352,16 +1352,17 @@
 ::new);%0A
+%0A
     @Ali
@@ -1372,237 +1372,83 @@
     
-@RecomputeFieldValue(kind = RecomputeFieldValue.Kind.Reset)%0A    private Map%3CSchema, Class%3E stringClassCache;%0A    @Alias%0A    @RecomputeFieldValue(kind = RecomputeFieldValue.Kind.Reset)%0A    private Map%3CClass, Constructor%3E stringCto
+private Target_org_apache_avro_generic_GenericDatumReader_ReaderCache reade
 rCac
@@ -1643,168 +1643,696 @@
 his.
-stringClassCache = new IdentityHashMap(
+data = data;%0A        this.creator = Thread.currentThread();%0A        this.readerCache = new Target_org_apache_avro_generic_GenericDatumReader_ReaderCache(this::findStringClass
 );%0A    
+%7D%0A%0A
     
-this.stringCtorCache = new HashMap();%0A        this.data = data;%0A        this.creator = Thread.currentThread();
+@Alias%0A    protected native Class findStringClass(Schema schema);%0A%0A%7D%0A%0A@TargetClass(className = %22org.apache.avro.generic.GenericDatumReader%22, innerClass = %22ReaderCache%22)%0Afinal class Target_org_apache_avro_generic_GenericDatumReader_ReaderCache %7B%0A    // Just provide access to %22org.apache.avro.generic.GenericDatumReader.ReaderCache%22%0A%0A    @Alias%0A    public Target_org_apache_avro_generic_GenericDatumReader_ReaderCache(%0A            @SuppressWarnings(%22unused%22) Function%3CSchema, Class%3E findStringClass) %7B
 %0A   

Fix Avro GenericDatumReader substitution