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

import com.squareup.javapoet.*;
import io.gitlab.jfronny.commons.StringFormatter;
import io.gitlab.jfronny.gson.compile.annotations.GPrefer;
import io.gitlab.jfronny.gson.compile.annotations.GSerializable;
import io.gitlab.jfronny.gson.compile.processor.adapter.Adapter;
import io.gitlab.jfronny.gson.compile.processor.adapter.Adapters;
import io.gitlab.jfronny.gson.compile.processor.core.AbstractProcessor2;
import io.gitlab.jfronny.gson.compile.processor.core.SupportedAnnotationTypes2;
import io.gitlab.jfronny.gson.compile.processor.core.value.ElementException;
import io.gitlab.jfronny.gson.compile.processor.core.value.ValueCreator;
import io.gitlab.jfronny.gson.compile.processor.gprocessor.*;
import io.gitlab.jfronny.gson.compile.processor.core.StringListComparator;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

@SupportedSourceVersion(SourceVersion.RELEASE_17)
@SupportedAnnotationTypes2({GSerializable.class})
@SupportedOptions({"gsonCompileNoReflect"})
public class GsonCompileProcessor extends AbstractProcessor2 {
    private Map<ClassName, TypeSpec.Builder> seen;
    private InstanceProcessor instanceProcessor;
    private StaticProcessor staticProcessor;

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

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnvironment) {
        Set<SerializableClass> toGenerate = 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, bld.builder, bld.configure, bld.generateAdapter, bld.isStatic, hasManifold));
                        }
                    } catch (ElementException e) {
                        e.printMessage(message);
                    }
                }
            }
        }
        // Do not allow mutation past this point, especially not from individual process tasks
        toGenerate = Set.copyOf(toGenerate);
        // Generate adapters
        for (SerializableClass toProcess : toGenerate) {
            try {
                process(toProcess, toGenerate);
            } 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));
                    }
                }
            }
        }
        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) 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());
        } else {
            if (toProcess.generateAdapter()) {
                processor.generateDelegatingAdapter(spec, classType, toProcess.generatedClassName());
            }
            processor.generateSerialisation(spec, toProcess, typeVariables, other);
        }

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