package io.gitlab.jfronny.commons.serialize.generator;

import com.squareup.javapoet.*;
import io.gitlab.jfronny.commons.StringFormatter;
import io.gitlab.jfronny.commons.serialize.databind.api.SerializeWithAdapter;
import io.gitlab.jfronny.commons.serialize.databind.api.SerializerFor;
import io.gitlab.jfronny.commons.serialize.generator.adapter.Adapter;
import io.gitlab.jfronny.commons.serialize.generator.adapter.Adapters;
import io.gitlab.jfronny.commons.serialize.generator.annotations.GPrefer;
import io.gitlab.jfronny.commons.serialize.generator.annotations.GSerializable;
import io.gitlab.jfronny.commons.serialize.generator.core.AbstractProcessor2;
import io.gitlab.jfronny.commons.serialize.generator.core.StringListComparator;
import io.gitlab.jfronny.commons.serialize.generator.core.SupportedAnnotationTypes2;
import io.gitlab.jfronny.commons.serialize.generator.core.value.ElementException;
import io.gitlab.jfronny.commons.serialize.generator.core.value.ValueCreator;
import io.gitlab.jfronny.commons.serialize.generator.gprocessor.GProcessor;
import io.gitlab.jfronny.commons.serialize.generator.gprocessor.InstanceProcessor;
import io.gitlab.jfronny.commons.serialize.generator.gprocessor.StaticProcessor;

import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedOptions;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import javax.tools.StandardLocation;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

@SupportedSourceVersion(SourceVersion.RELEASE_17)
@SupportedAnnotationTypes2({GSerializable.class, SerializeWithAdapter.class, SerializerFor.class})
@SupportedOptions({"serializeProcessorNoReflect", "serializeProcessorDisableSafe"})
public class SerializeGeneratorProcessor extends AbstractProcessor2 {
    private Map<ClassName, TypeSpec.Builder> seen;
    private InstanceProcessor instanceProcessor;
    private StaticProcessor staticProcessor;
    private List<ClassName> generatedAdapters;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        ValueCreator.preferAnnotation = GPrefer.class;
        seen = new LinkedHashMap<>();
        generatedAdapters = new ArrayList<>();
        for (Adapter adapter : Adapters.ADAPTERS) {
            adapter.init(processingEnv);
        }
        boolean disableSafe = options.containsKey("serializeProcessorDisableSafe");
        instanceProcessor = new InstanceProcessor(valueCreator, message, hasManifold, disableSafe);
        staticProcessor = new StaticProcessor(valueCreator, message, hasManifold, disableSafe);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnvironment) {
        Set<SerializableClass> toGenerate = new LinkedHashSet<>();
        Set<AdapterRef> adapterRefs = new LinkedHashSet<>();
        // Gather all serializable types
        for (TypeElement annotation : annotations) {
            for (Element element : roundEnvironment.getElementsAnnotatedWith(annotation)) {
                for (AnnotationMirror mirror : element.getAnnotationMirrors()) {
                    try {
                        if (mirror.getAnnotationType().toString().equals(GSerializable.class.getCanonicalName())) {
                            var bld = new Object() {
                                TypeMirror with = null;
                                TypeMirror builder = null;
                                TypeMirror configure = null;
                                Boolean generateAdapter = null;
                                Boolean isStatic = null;
                            };
                            elements.getElementValuesWithDefaults(mirror).forEach((executableElement, value) -> {
                                String name = executableElement.getSimpleName().toString();
                                switch (name) {
                                    case "with" -> {
                                        if (bld.with != null) throw new IllegalArgumentException("Duplicate annotation parameter: with");
                                        bld.with = (TypeMirror) value.getValue();
                                    }
                                    case "builder" -> {
                                        if (bld.builder != null) throw new IllegalArgumentException("Duplicate annotation parameter: builder");
                                        bld.builder = (TypeMirror) value.getValue();
                                    }
                                    case "configure" -> {
                                        if (bld.configure != null) throw new IllegalArgumentException("Duplicate annotation parameter: configure");
                                        bld.configure = (TypeMirror) value.getValue();
                                    }
                                    case "generateAdapter" -> {
                                        if (bld.generateAdapter != null) throw new IllegalArgumentException("Duplicate annotation parameter: generateAdapter");
                                        bld.generateAdapter = (Boolean) value.getValue();
                                    }
                                    case "isStatic" -> {
                                        if (bld.isStatic != null) throw new IllegalArgumentException("Duplicate annotation parameter: isStatic");
                                        bld.isStatic = (Boolean) value.getValue();
                                    }
                                    default -> throw new IllegalArgumentException("Unexpected annotation parameter: " + name);
                                }
                            });
                            if (bld.with == null) throw new IllegalArgumentException("Missing annotation parameter: with");
                            if (bld.builder == null) throw new IllegalArgumentException("Missing annotation parameter: builder");
                            if (bld.configure == null) throw new IllegalArgumentException("Missing annotation parameter: configure");
                            if (bld.generateAdapter == null) throw new IllegalArgumentException("Missing annotation parameter: generateAdapter");
                            if (bld.isStatic == null) throw new IllegalArgumentException("Missing annotation parameter: isStatic");

                            toGenerate.add(SerializableClass.of((TypeElement) element, bld.with, true, bld.builder, bld.configure, bld.generateAdapter, bld.isStatic, hasManifold));
                        } else if (mirror.getAnnotationType().toString().equals(SerializeWithAdapter.class.getCanonicalName())) {
                            var bld = new Object() {
                                TypeMirror adapter = null;
                                Boolean nullSafe = null;
                            };
                            elements.getElementValuesWithDefaults(mirror).forEach((executableElement, value) -> {
                                String name = executableElement.getSimpleName().toString();
                                switch (name) {
                                    case "adapter" -> {
                                        if (bld.adapter != null) throw new IllegalArgumentException("Duplicate annotation parameter: adapter");
                                        if (value.getValue() instanceof TypeMirror tm) bld.adapter = tm;
                                        else throw new IllegalArgumentException("Unexpected value type: " + value.getValue().getClass() + " (with value" + value.getValue() + ")");
                                    }
                                    case "nullSafe" -> {
                                        if (bld.nullSafe != null) throw new IllegalArgumentException("Duplicate annotation parameter: nullSafe");
                                        bld.nullSafe = (Boolean) value.getValue();
                                    }
                                    default -> throw new IllegalArgumentException("Unexpected annotation parameter: " + name);
                                }
                            });
                            if (bld.adapter == null) throw new IllegalArgumentException("Missing annotation parameter: adapter");
                            if (bld.nullSafe == null) throw new IllegalArgumentException("Missing annotation parameter: nullSafe");

                            toGenerate.add(SerializableClass.of((TypeElement) element, bld.adapter, false, null, null, false, false, hasManifold));
                        } else if (mirror.getAnnotationType().toString().equals(SerializerFor.class.getCanonicalName())) {
                            var bld = new Object() {
                                TypeMirror[] targets = null;
                                Boolean nullSafe = null;
                            };
                            elements.getElementValuesWithDefaults(mirror).forEach((executableElement, value) -> {
                                String name = executableElement.getSimpleName().toString();
                                switch (name) {
                                    case "targets" -> {
                                        if (bld.targets != null) throw new IllegalArgumentException("Duplicate annotation parameter: targets");
                                        bld.targets = ((List<?>) value.getValue()).stream()
                                                .map(o -> (AnnotationValue) o)
                                                .map(o -> (TypeMirror) o.getValue())
                                                .toArray(TypeMirror[]::new);
                                    }
                                    case "nullSafe" -> {
                                        if (bld.nullSafe != null) throw new IllegalArgumentException("Duplicate annotation parameter: nullSafe");
                                        bld.nullSafe = (Boolean) value.getValue();
                                    }
                                    case "hierarchical" -> {
                                        // not relevant for this processor
                                    }
                                    default -> throw new IllegalArgumentException("Unexpected annotation parameter: " + name);
                                }
                            });
                            if (bld.targets == null) throw new IllegalArgumentException("Missing annotation parameter: targets");
                            if (bld.nullSafe == null) throw new IllegalArgumentException("Missing annotation parameter: nullSafe");
                            adapterRefs.add(new AdapterRef((TypeElement) element, bld.targets, bld.nullSafe));
                        } else {
                            throw new IllegalArgumentException("Unexpected annotation: " + mirror.getAnnotationType());
                        }
                    } catch (ElementException e) {
                        e.printMessage(message);
                    }
                }
            }
        }
        // Do not allow mutation past this point, especially not from individual process tasks
        toGenerate = Set.copyOf(toGenerate);
        adapterRefs = Set.copyOf(adapterRefs);
        // Generate adapters
        for (SerializableClass toProcess : toGenerate) {
            try {
                process(toProcess, toGenerate, adapterRefs);
            } catch (ElementException e) {
                e.printMessage(message);
            }
        }
        for (var entry : seen.keySet().stream().collect(Collectors.groupingBy(ClassName::packageName)).entrySet()) {
            Map<List<String>, TypeSpec.Builder> known = entry.getValue().stream()
                    .collect(Collectors.toMap(
                            ClassName::simpleNames,
                            seen::get,
                            (u, v) -> u,
                            () -> new TreeMap<>(StringListComparator.INSTANCE.reversed())
                    ));
            // Generate additional parent classes
            for (List<String> klazz : known.keySet().stream().toList()) {
                List<String> current = new LinkedList<>();
                for (String s : klazz) {
                    current.add(s);
                    TypeSpec.Builder builder = find(known, current);
                    if (builder == null) {
                        builder = TypeSpec.classBuilder(s).addModifiers(Modifier.PUBLIC);
                        if (current.size() == 1 && hasManifold) builder.addAnnotation(Cl.MANIFOLD_EXTENSION);
                        known.put(List.copyOf(current), builder);
                    }
                    if (current.size() > 1) builder.addModifiers(Modifier.STATIC);
                }
            }
            // Add to parent class
            for (var entry1 : known.entrySet()) {
                if (entry1.getKey().size() == 1) continue;
                find(known, entry1.getKey().subList(0, entry1.getKey().size() - 1)).addType(entry1.getValue().build());
            }
            // Print
//            System.out.println("Got " + known.size() + " classes");
//            for (var entry1 : known.entrySet()) {
//                System.out.println("Class " + entry.key + '.' + String.join(".", entry1.key));
//                for (TypeSpec typeSpec : entry1.value.typeSpecs) {
//                    System.out.println("- " + typeSpec.name);
//                }
//            }
            // Write top-level classes
            for (var entry1 : known.entrySet()) {
                if (entry1.getKey().size() == 1) {
                    JavaFile javaFile = JavaFile.builder(entry.getKey(), entry1.getValue().build())
                            .skipJavaLangImports(true)
                            .indent("    ")
                            .build();
                    try {
                        javaFile.writeTo(filer);
                    } catch (IOException e) {
                        message.printMessage(Diagnostic.Kind.ERROR, "Could not write source: " + StringFormatter.toString(e));
                    }
                }
            }
        }
        if (!generatedAdapters.isEmpty()) {
            try {
                filer.createResource(StandardLocation.SOURCE_OUTPUT, "", "META-INF/services/io.gitlab.jfronny.commons.serialize.databind.api.TypeAdapter")
                        .openWriter()
                        .append(Stream.concat(generatedAdapters.stream(), adapterRefs.stream().map(AdapterRef::in).map(ClassName::get))
                                .map(ClassName::canonicalName)
                                .collect(Collectors.joining("\n")))
                        .close();
            } catch (IOException e) {
                message.printMessage(Diagnostic.Kind.ERROR, "Could not write service loader file: " + StringFormatter.toString(e));
            }
            generatedAdapters.clear();
        }
        seen.clear();
        return false;
    }

    private <K, V> V find(Map<List<K>, V> map, List<K> key) {
        for (var entry : map.entrySet()) if (entry.getKey().equals(key)) return entry.getValue();
        return null;
    }

    private void process(SerializableClass toProcess, Set<SerializableClass> other, Set<AdapterRef> refs) throws ElementException {
        if (seen.containsKey(toProcess.generatedClassName())) return; // Don't process the same class more than once

        TypeName classType = toProcess.getTypeName();
        List<TypeVariableName> typeVariables = new ArrayList<>();
        if (!toProcess.isStatic() && classType instanceof ParameterizedTypeName type) {
            for (TypeName argument : type.typeArguments) {
                typeVariables.add(TypeVariableName.get(argument.toString()));
            }
        }

        TypeSpec.Builder spec = TypeSpec.classBuilder(toProcess.generatedClassName().simpleName())
                .addOriginatingElement(toProcess.classElement())
                .addModifiers(Modifier.PUBLIC);
//                .addTypeVariables(typeVariables);

        seen.put(toProcess.generatedClassName(), spec);

        GProcessor processor = toProcess.isStatic() ? staticProcessor : instanceProcessor;

        if (toProcess.adapter() != null) {
            processor.generateDelegateToAdapter(spec, classType, toProcess.adapter(), toProcess.adapterIsStatic());
        } else {
            if (toProcess.generateAdapter()) {
                generatedAdapters.add(processor.generateDelegatingAdapter(spec, classType, toProcess.generatedClassName()));
            }
            processor.generateSerialisation(spec, toProcess, typeVariables, other, refs); //TODO fix type annotations
        }

        processor.generateAuxiliary(spec, classType, toProcess.configure());
    }
}