package io.gitlab.jfronny.libjf.config.plugin;

import com.squareup.javapoet.*;
import io.gitlab.jfronny.commons.StringFormatter;
import io.gitlab.jfronny.commons.throwable.Coerce;
import io.gitlab.jfronny.gson.compile.processor.core.*;
import io.gitlab.jfronny.gson.compile.processor.core.value.*;
import io.gitlab.jfronny.gson.compile.processor.core.value.Properties;
import io.gitlab.jfronny.gson.reflect.TypeToken;
import io.gitlab.jfronny.libjf.config.api.v1.*;
import io.gitlab.jfronny.libjf.config.api.v1.dsl.DSL;
import io.gitlab.jfronny.libjf.config.api.v1.type.Type;

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

@SupportedSourceVersion(SourceVersion.RELEASE_17)
@SupportedAnnotationTypes2({JfConfig.class})
@SupportedOptions({"modId"})
public class ConfigProcessor extends AbstractProcessor2 {
    private Map<ClassName, TypeSpec.Builder> seen;
    private String modId = "null";

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        seen = new LinkedHashMap<>();
        if (!options.containsKey("modId")) message.printError("Lacking modId: can't set up config");
        else modId = options.get("modId");
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnvironment) {
        Set<ConfigClass> toGenerate = new LinkedHashSet<>();
        // Gather all serializable types
        annotations.stream().flatMap(s -> roundEnvironment.getElementsAnnotatedWith(s).stream()).forEach(element -> {
            JfConfig jc = element.getAnnotation(JfConfig.class);
            if (jc != null) {
                toGenerate.add(ConfigClass.of((TypeElement) element, jc.referencedConfigs(), hasManifold));
            }
        });
        // Generate adapters
        toGenerate.forEach(Coerce.<ConfigClass, ElementException>consumer(toProcess -> process(toProcess, toGenerate)).addHandler(ElementException.class, 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(ConfigClass toProcess, Set<ConfigClass> other) throws ElementException {
        if (seen.containsKey(toProcess.generatedClassName())) return; // Don't process the same class more than once

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

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

        CodeBlock.Builder code = CodeBlock.builder();
        code.addStatement("$T.getInstance().migrateFiles($S)", ConfigHolder.class, modId);
        code.add("INSTANCE = $T.create($S).register(builder0 -> builder0", DSL.class, modId).indent();
        for (String s : toProcess.referencedConfigs()) code.add("\n.referenceConfig($S)", s);
        process(toProcess.classElement(), code, new AtomicInteger(0));
        code.unindent().add("\n);\n");

        spec.addField(FieldSpec.builder(double.class, "Infinity", Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL).initializer("$T.POSITIVE_INFINITY", Double.class).build());
        spec.addStaticBlock(code.build());
        spec.addField(FieldSpec.builder(ConfigInstance.class, "INSTANCE", Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL).build());
        spec.addMethod(MethodSpec.methodBuilder("write").addModifiers(Modifier.PUBLIC, Modifier.STATIC).addCode("INSTANCE.write();").build());
        spec.addMethod(MethodSpec.methodBuilder("load").addModifiers(Modifier.PUBLIC, Modifier.STATIC).addCode("INSTANCE.load();").build());
    }

    private void process(TypeElement source, CodeBlock.Builder code, AtomicInteger i) {
        for (TypeElement klazz : ElementFilter.typesIn(source.getEnclosedElements())) {
            Category v = klazz.getAnnotation(Category.class);
            if (v != null) {
                String name = klazz.getSimpleName().toString();
                name = Character.toLowerCase(name.charAt(0)) + name.substring(1);
                code.add("\n.category($S, builder$L -> builder$L", name, i.incrementAndGet(), i.get()).indent();
                for (String s : v.referencedConfigs()) code.add("\n.referencedConfigs($S)", s);
                process(klazz, code, i);
                code.unindent().add("\n)");
            }
        }
        for (VariableElement field : ElementFilter.fieldsIn(source.getEnclosedElements())) {
            if (field.getModifiers().containsAll(Set.of(Modifier.PUBLIC, Modifier.STATIC))) {
                Entry e = field.getAnnotation(Entry.class);
                if (e == null) continue;
                String name = field.getSimpleName().toString();
                TypeMirror tm = unbox(field.asType());
                code.add("\n.value($S, $T.$L, ", name, source, name);
                DeclaredType declared = TypeHelper.asDeclaredType(tm);
                if (declared != null && declared.asElement().getKind() == ElementKind.ENUM) {
                    code.add("$T.class, () -> $T.$L, value -> $T.$L = value", declared, source, name, source, name);
                } else if (tm.toString().equals(String.class.getCanonicalName())) {
                    code.add("() -> $T.$L, value -> $T.$L = value", source, name, source, name);
                } else switch (tm.getKind()) {
                    case BOOLEAN -> code.add("() -> $T.$L, value -> $T.$L = value", source, name, source, name);
                    case BYTE -> code.add("$L, $L, () -> (int) $T.$L, value -> $T.$L = (byte) (int) value", e.min(), e.max(), source, name, source, name);
                    case SHORT -> code.add("$L, $L, () -> (int) $T.$L, value -> $T.$L = (short) (int) value", e.min(), e.max(), source, name, source, name);
                    case CHAR -> code.add("$L, $L, () -> (int) $T.$L, value -> $T.$L = (char) (int) value", e.min(), e.max(), source, name, source, name);
                    case INT, LONG, FLOAT, DOUBLE -> code.add("$L, $L, () -> $T.$L, value -> $T.$L = value", e.min(), e.max(), source, name, source, name);
                    default -> {
                        code.add("$L, $L, ", e.min(), e.max());
                        if (tm instanceof DeclaredType dt && !dt.getTypeArguments().isEmpty()) {
                            code.add("$T.ofClass(new $T<$T>() {}.getType())", Type.class, TypeToken.class, tm);
                        } else code.add("$T.class", tm);
                        code.add(", $L, () -> $T.$L, value -> $T.$L = value", e.width(), source, name, source, name);
                    }
                }
                code.add(")");
            }
        }
        for (ExecutableElement method : ElementFilter.methodsIn(source.getEnclosedElements())) {
            if (method.getModifiers().containsAll(Set.of(Modifier.PUBLIC, Modifier.STATIC))) {
                String name = method.getSimpleName().toString();
                Preset p = method.getAnnotation(Preset.class);
                if (p != null) code.add("\n.addPreset($S, $T::$L)", name, source, name);
                Verifier v = method.getAnnotation(Verifier.class);
                if (v != null) code.add("\n.addVerifier($T::$L)", source, name);
            }
        }
    }

    private TypeMirror unbox(TypeMirror type) {
        try {
            return types.unboxedType(type);
        } catch (IllegalArgumentException e) {
            return type;
        }
    }
}