package cn.com.v2.common.interceptor;

import cn.com.v2.util.DataSourceUtil;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;

@Aspect
@Component
public class DataSourceAspect {

    // 处理方法和类上同时存在@DS注解，以方法上的为准
    private final ThreadLocal<Boolean> methodHasAnnotationProcessed = ThreadLocal.withInitial(() -> false);

    @Before("@within(DS) || @annotation(DS)")
    public void handleDataSourceChange(JoinPoint joinPoint) {
        if (!methodHasAnnotationProcessed.get()) {
            MethodSignature signature = (MethodSignature) joinPoint.getSignature();
            Method method = signature.getMethod();
            DS methodAnnotation = AnnotationUtils.findAnnotation(method, DS.class);

            if (methodAnnotation != null && StringUtils.hasText(methodAnnotation.dbName())) {
                DataSourceUtil.setDB(methodAnnotation.dbName());
                methodHasAnnotationProcessed.set(true);
            } else {
                // 方法上没找到就找下类上是否有注解
                Class<?> declaringClass = method.getDeclaringClass();
                DS classAnnotation = declaringClass.getAnnotation(DS.class);
                if (classAnnotation != null && StringUtils.hasText(classAnnotation.dbName())) {
                    DataSourceUtil.setDB(classAnnotation.dbName());
                    methodHasAnnotationProcessed.set(true);
                }
            }
        }
    }

    // 在每个线程执行结束时清理ThreadLocal变量，避免累积问题，表达式中com.example记得换成自己的包路径
    @AfterReturning("execution(* cn.com..*Controller.*(..))")
    public void resetControllerAnnotationProcessed() {
        methodHasAnnotationProcessed.remove();
    }
}
