package org.hiforce.lattice.spi.annotation;

import com.google.common.base.Optional;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;
import com.google.common.collect.Sets;
import org.hiforce.lattice.utils.ServicesFileUtils;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Filer;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.tools.Diagnostic;
import javax.tools.FileObject;
import javax.tools.StandardLocation;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.annotation.Annotation;
import java.util.HashSet;
import java.util.Set;
import java.util.SortedSet;

import static com.google.auto.common.MoreElements.getAnnotationMirror;

/**
 * Lattice框架的注解处理器基类，用于在编译期处理自定义注解
 * 实现了Java SPI机制的自动注册功能，将标注的类自动写入META-INF/services配置文件
 * 
 * 工作原理：
 * 1. 扫描源码中的特定注解（如@Extension、@Business等）
 * 2. 收集被注解标记的类信息
 * 3. 在编译完成时生成META-INF/services文件，实现SPI自动注册
 * 
 * @author Rocky Yu
 * @since 2022/9/15
 */
@SuppressWarnings("all")
public abstract class LatticeAnnotationProcessor extends AbstractProcessor {

    /**
     * 获取服务接口的Class对象
     * 子类需要实现此方法，指定要注册到SPI的接口类型
     * @return 服务接口的Class对象
     */
    public abstract Class<?> getServiceInterfaceClass();

    /**
     * 日志输出方法，仅在debug模式下输出
     * @param msg 日志消息
     */
    private void log(String msg) {
        if (processingEnv.getOptions().containsKey("debug")) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.NOTE, msg);
        }
    }

    /**
     * 输出错误信息到编译器
     * @param msg 错误消息
     * @param element 出错的元素
     * @param annotation 出错的注解
     */
    private void error(String msg, Element element, AnnotationMirror annotation) {
        processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, msg, element, annotation);
    }

    /**
     * 输出致命错误，会导致编译失败
     * @param msg 错误消息
     */
    private void fatalError(String msg) {
        processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "FATAL ERROR: " + msg);
    }

    // 错误提示常量：没有为元素提供服务接口
    public static final String MISSING_SERVICES_ERROR = "No service interfaces provided for element!";

    /**
     * 服务提供者映射表
     * Key: 服务接口的全限定类名（如：org.hiforce.lattice.spi.annotation.ExtensionAnnotationParser）
     * Value: 实现类的全限定类名集合（如：org.hiforce.lattice.annotation.parser.DefaultExtensionAnnotationParser）
     * 用于收集所有需要注册到SPI的实现类
     */
    private final Multimap<String, String> providers = HashMultimap.create();

    /**
     * 注解处理的入口方法，由编译器调用
     * @param annotations 本轮处理的注解类型集合
     * @param roundEnv 当前处理轮次的环境信息
     * @return true表示这些注解已被处理，其他处理器不应再处理它们
     */
    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        try {
            // 调用实际的处理逻辑
            return processImpl(annotations, roundEnv);
        } catch (Exception e) {
            // 捕获所有异常，防止异常传播到编译器导致编译中断
            StringWriter writer = new StringWriter();
            e.printStackTrace(new PrintWriter(writer));
            fatalError(writer.toString());
            return true;
        }
    }

    /**
     * 注解处理的核心实现方法
     * 注解处理分为多轮进行，最后一轮用于生成配置文件
     * @param annotations 本轮处理的注解类型集合
     * @param roundEnv 当前处理轮次的环境信息
     * @return 总是返回true
     */
    private boolean processImpl(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        if (roundEnv.processingOver()) {
            // 最后一轮处理：生成META-INF/services配置文件
            generateConfigFiles();
        } else {
            // 常规处理轮次：扫描和收集被注解标记的类
            processAnnotations(annotations, roundEnv);
        }

        return true;
    }

    /**
     * 获取要处理的注解类型
     * 子类需要实现此方法，指定要处理哪个注解（如@Extension、@Business等）
     * @return 注解的Class对象
     */
    public abstract Class<? extends Annotation> getProcessAnnotationClass();

    /**
     * 返回此处理器支持的注解类型集合
     * @return 包含注解全限定类名的不可变集合
     */
    @Override
    public final ImmutableSet<String> getSupportedAnnotationTypes() {
        return ImmutableSet.of(getProcessAnnotationClass().getName());
    }

    /**
     * 处理注解标记的元素，收集需要注册的服务实现类
     * 支持两种注解位置：
     * 1. 类级别注解：直接标注在类上
     * 2. 方法级别注解：标注在方法上，注册方法所在的类
     * 
     * @param annotations 注解类型集合
     * @param roundEnv 当前轮次环境
     */
    @SuppressWarnings("all")
    private void processAnnotations(Set<? extends TypeElement> annotations,
                                    RoundEnvironment roundEnv) {

        // 获取所有被目标注解标记的元素（类或方法）
        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(getProcessAnnotationClass());

        log(annotations.toString());
        log(elements.toString());

        // 遍历所有被注解标记的元素
        for (Element e : elements) {
            TypeElement providerImplementer = null;
            
            // 情况1：注解标注在类上
            if (e instanceof TypeElement) {
                providerImplementer = (TypeElement) e;
            }
            
            // 情况2：注解标注在方法上，获取方法所在的类
            if (e.getKind() == ElementKind.METHOD && e instanceof ExecutableElement) {
                ExecutableElement executableElement = (ExecutableElement) e;
                Element element = executableElement.getEnclosingElement();
                // 确保外围元素是类类型
                if (element.getKind() == ElementKind.CLASS && element instanceof TypeElement) {
                    providerImplementer = (TypeElement) (executableElement.getEnclosingElement());
                }
            }
            
            // 如果无法确定实现类，跳过此元素
            if (providerImplementer == null) {
                continue;
            }
            
            // 获取注解镜像对象，用于错误报告
            Optional<AnnotationMirror> optional = getAnnotationMirror(e, getProcessAnnotationClass());
            if (optional.isPresent()) {
                AnnotationMirror annotationMirror = optional.get();
                // 获取服务接口类
                Class<?> serviceClass = getServiceInterfaceClass();
                if (null == serviceClass) {
                    error(MISSING_SERVICES_ERROR, e, annotationMirror);
                    continue;
                }
                // 将实现类注册到providers映射表中
                // Key: 服务接口全限定名, Value: 实现类全限定名
                providers.put(serviceClass.getName(), getBinaryName(providerImplementer));
            }
        }
    }

    /**
     * 生成META-INF/services配置文件，实现Java SPI机制的自动注册
     * 
     * 工作流程：
     * 1. 读取已存在的服务配置文件（如果有）
     * 2. 合并新发现的服务实现类
     * 3. 写入更新后的配置文件
     * 
     * 生成的文件格式：META-INF/services/接口全限定名
     * 文件内容：每行一个实现类的全限定名
     */
    private void generateConfigFiles() {
        // 获取文件操作对象，用于创建和读取资源文件
        Filer filer = processingEnv.getFiler();

        // 遍历所有收集到的服务接口
        for (String providerInterface : providers.keySet()) {
            // 构建SPI配置文件路径：META-INF/services/接口全限定名
            String resourceFile = "META-INF/services/" + providerInterface;
            log("Working on resource file: " + resourceFile);
            try {
                // 使用TreeSet保证实现类列表有序
                SortedSet<String> allServices = Sets.newTreeSet();
                try {
                    // 尝试读取已存在的服务配置文件
                    // 这样可以支持增量编译，保留之前注册的服务
                    FileObject existingFile = filer.getResource(StandardLocation.CLASS_OUTPUT, "",
                            resourceFile);
                    log("Looking for existing resource file at " + existingFile.toUri());
                    Set<String> oldServices = ServicesFileUtils.readServiceFile(existingFile.openInputStream());
                    log("Existing service entries: " + oldServices);
                    allServices.addAll(oldServices);
                } catch (IOException e) {
                    // 文件不存在时会抛出IOException
                    // 这是正常情况，说明是首次编译或配置文件尚未创建
                    log("Resource file did not already exist.");
                }

                // 获取本次编译新发现的服务实现类
                Set<String> newServices = new HashSet<String>(providers.get(providerInterface));
                
                // 如果所有新服务都已存在，无需重新生成文件
                if (allServices.containsAll(newServices)) {
                    log("No new service entries being added.");
                    return;
                }

                // 合并新旧服务实现类
                allServices.addAll(newServices);
                log("New service file contents: " + allServices);
                
                // 创建或覆盖SPI配置文件
                FileObject fileObject = filer.createResource(StandardLocation.CLASS_OUTPUT, "",
                        resourceFile);
                OutputStream out = fileObject.openOutputStream();
                // 将所有服务实现类写入文件，每行一个类名
                ServicesFileUtils.writeServiceFile(allServices, out);
                out.close();
                log("Wrote to: " + fileObject.toUri());
            } catch (IOException e) {
                fatalError("Unable to createPluginConfig " + resourceFile + ", " + e);
                return;
            }
        }
    }

    /**
     * 获取类型元素的二进制名称（全限定类名）
     * 处理内部类时使用$分隔符，例如：com.google.Foo$Bar
     * 而不是源码中的点分隔：com.google.Foo.Bar
     * 
     * @param element 类型元素
     * @return 二进制格式的全限定类名
     */
    private String getBinaryName(TypeElement element) {
        return getBinaryNameImpl(element, element.getSimpleName().toString());
    }

    /**
     * 递归构建二进制类名的实现方法
     * 通过向上遍历外围元素，构建完整的类名
     * 
     * @param element 当前类型元素
     * @param className 累积的类名（从内向外构建）
     * @return 完整的二进制类名
     */
    private String getBinaryNameImpl(TypeElement element, String className) {
        // 获取外围元素（可能是包或外部类）
        Element enclosingElement = element.getEnclosingElement();

        // 情况1：外围元素是包，说明已到达顶层类
        if (enclosingElement instanceof PackageElement) {
            PackageElement pkg = (PackageElement) enclosingElement;
            // 处理默认包（无名包）的情况
            if (pkg.isUnnamed()) {
                return className;
            }
            // 返回：包名.类名
            return pkg.getQualifiedName() + "." + className;
        }

        // 情况2：外围元素是类，说明当前类是内部类
        TypeElement typeElement = (TypeElement) enclosingElement;
        // 递归向上，使用$连接外部类和内部类
        // 例如：Outer$Inner$DeepInner
        return getBinaryNameImpl(typeElement, typeElement.getSimpleName() + "$" + className);
    }

    /**
     * 返回此注解处理器支持的Java源代码版本
     * @return Java 8版本
     */
    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.RELEASE_8;
    }
}