package cn.axzo.workflow.generate;

import com.google.auto.service.AutoService;
import com.squareup.javapoet.MethodSpec;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.Processor;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.ElementFilter;
import javax.tools.Diagnostic;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;

@AutoService(Processor.class)
public class WorkflowExposeApiGenProcessor extends AbstractProcessor {

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        List<GenCodeInfo> genCodeInfos = new ArrayList<>();
        annotations.forEach(an -> {
            Set<? extends Element> typeElements = roundEnv.getElementsAnnotatedWith(an);
            Set<TypeElement> types = ElementFilter.typesIn(typeElements);
            for (TypeElement typeElement : types) {
                CodeGenProcessor codeGenProcessor = CodeGenProcessorRegistry.find(an.getQualifiedName().toString());
                try {
                    GenCodeInfo generate = codeGenProcessor.generate(typeElement, roundEnv);
                    genCodeInfos.add(generate);
                } catch (Exception e) {
                    ProcessingEnvironmentHolder.getEnvironment().getMessager().printMessage(Diagnostic.Kind.ERROR, "代码生成异常:" + e.getMessage());
                }
            }
        });
        if (!genCodeInfos.isEmpty()) {
            Map<Triplet<String, String, String>, List<GenCodeInfo>> apiGenMap = new HashMap<>();
            //进行分组
            for (GenCodeInfo genCodeInfo : genCodeInfos) {
                Triplet<String, String, String> triplet = Triplet.of(genCodeInfo.getSourcePath(), genCodeInfo.getPackageName(), genCodeInfo.getClassName());
                List<GenCodeInfo> codeInfos = apiGenMap.computeIfAbsent(triplet, k -> new ArrayList<>());
                codeInfos.add(genCodeInfo);
            }
            for (Map.Entry<Triplet<String, String, String>, List<GenCodeInfo>> entry : apiGenMap.entrySet()) {
                Triplet<String, String, String> triplet = entry.getKey();
                List<MethodSpec> methodSpecs = new ArrayList<>();
                List<GenCodeInfo> infos = entry.getValue();
                for (GenCodeInfo genCodeInfo : infos) {
                    methodSpecs.addAll(genCodeInfo.getTypeSpec().methodSpecs);
                }
                GenJavaSourceFileUtils.genJavaSourceFile(triplet.getFirst(), triplet.getSecond(), triplet.getThird(), methodSpecs);
            }
        }
        return false;
    }

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        ProcessingEnvironmentHolder.setEnvironment(processingEnv);
        CodeGenProcessorRegistry.initProcessors();
    }

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return CodeGenProcessorRegistry.getSupportedAnnotations();
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    public static final class Triplet<T, U, V> {
        private final T first;
        private final U second;
        private final V third;

        public Triplet(T first, U second, V third) {
            this.first = first;
            this.second = second;
            this.third = third;
        }

        public T getFirst() {
            return first;
        }

        public U getSecond() {
            return second;
        }

        public V getThird() {
            return third;
        }

        @Override
        public String toString() {
            return "Triplet{" +
                    "first=" + first +
                    ", second=" + second +
                    ", third=" + third +
                    '}';
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Triplet<?, ?, ?> triplet = (Triplet<?, ?, ?>) o;
            return Objects.equals(first, triplet.first) && Objects.equals(second, triplet.second) && Objects.equals(third, triplet.third);
        }

        @Override
        public int hashCode() {
            return Objects.hash(first, second, third);
        }

        public static <T, U, V> Triplet<T, U, V> of(T first, U second, V third) {
            return new Triplet<>(first, second, third);
        }
    }
}
