package com.centreproject.framework.config;

import com.centreproject.common.core.domain.model.LoginUser;
import com.centreproject.framework.web.service.TokenService;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.aot.hint.annotation.RegisterReflectionForBinding;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.util.concurrent.Executor;

@Configuration
@EnableAsync
@RegisterReflectionForBinding(SecurityContext.class) // 确保AOT模式下能处理安全上下文
public class AsyncConfig implements AsyncConfigurer {

    private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AsyncConfig.class);

    @Autowired
    private TokenService tokenService;

    @Override
    public Executor getAsyncExecutor() {
        if (Runtime.version().feature() >= 21) {
            return createVirtualThreadExecutor();
        }
        return createThreadPoolExecutor();
    }

    /**
     * 创建虚拟线程执行器（Spring Boot 3.5.4适配方案）
     */
    private Executor createVirtualThreadExecutor() {
        return command -> {
            // 1. 捕获当前请求上下文
            RequestAttributes attributes = RequestContextHolder.getRequestAttributes();
            LoginUser loginUser = tokenService.getLoginUser(
                    ((ServletRequestAttributes) attributes).getRequest()
            );

            // 2. 创建安全任务
            Runnable securedTask = () -> {
                try {
                    if (loginUser != null) {
                        SecurityContext context = SecurityContextHolder.createEmptyContext();
                        context.setAuthentication(
                                new UsernamePasswordAuthenticationToken(
                                        loginUser, null, loginUser.getAuthorities()
                                )
                        );
                        SecurityContextHolder.setContext(context);
                    }
                    command.run();
                } finally {
                    SecurityContextHolder.clearContext();
                }
            };

            // 3. 使用虚拟线程执行
            Thread.ofVirtual()
                    .name("VT-Async-")
                    .start(securedTask);
        };
    }

    /**
     * 传统线程池（兼容方案）
     */
    private Executor createThreadPoolExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(10);
        executor.setThreadNamePrefix("Pool-Async-");
        executor.setTaskDecorator(new Jdk21SecurityDecorator(tokenService));
        executor.initialize();
        return executor;
    }

    /**
     * 安全上下文装饰器（JDK 21完整实现）
     */
    static class Jdk21SecurityDecorator implements TaskDecorator {
        private final TokenService tokenService;

        public Jdk21SecurityDecorator(TokenService tokenService) {
            this.tokenService = tokenService;
        }

        @Override
        public Runnable decorate(Runnable runnable) {
            // 1. 捕获当前请求上下文（主线程）
            RequestAttributes attributes = RequestContextHolder.getRequestAttributes();
            HttpServletRequest request = attributes != null ?
                    ((ServletRequestAttributes) attributes).getRequest() : null;

            // 2. 提前解析用户信息（避免异步线程中访问Request）
            LoginUser loginUser = null;
            if (request != null) {
                try {
                    loginUser = tokenService.getLoginUser(request);
                } catch (Exception e) {
                    logger.error("异步任务安全上下文解析失败", e);
                }
            }

            // 3. 返回带安全上下文的Runnable
            LoginUser finalLoginUser = loginUser;
            return () -> {
                try {
                    // 4. 在新线程中重建安全上下文
                    if (finalLoginUser != null) {
                        SecurityContext context = SecurityContextHolder.createEmptyContext();
                        context.setAuthentication(
                                new UsernamePasswordAuthenticationToken(
                                        finalLoginUser,
                                        null,
                                        finalLoginUser.getAuthorities()
                                )
                        );
                        SecurityContextHolder.setContext(context);
                    }

                    // 5. 执行业务逻辑
                    runnable.run();
                } catch (Exception e) {
                    logger.error("异步任务执行异常", e);
                    throw e;
                } finally {
                    // 6. 必须清理线程状态
                    SecurityContextHolder.clearContext();
                    // 针对虚拟线程的特殊清理（防止内存泄漏）
                    if (Thread.currentThread().isVirtual()) {
                        RequestContextHolder.resetRequestAttributes();
                    }
                }
            };
        }
    }
}