@Data
@TableName("users")
public class User {
    @TableId(type = IdType.AUTO)
    private Long id;
    private String username;
    private String password;
    private Boolean enabled;
    private Boolean mfaEnabled;
    private String mfaSecret;
    private LocalDateTime createdAt;
    private LocalDateTime updatedAt;
}

@Mapper
public interface UserMapper extends BaseMapper<User> {
    // 如果需要自定义 SQL，可在此处声明
}

public class TotpUtils {
    private static final GoogleAuthenticator gAuth = new GoogleAuthenticator();

    /**
     * 为用户生成一个新的 TOTP 密钥（Base32 编码格式）
     *
     * @return Base32 编码的密钥
     */
    public static String generateSecretKey() {
        GoogleAuthenticatorKey key = gAuth.createCredentials();
        return key.getKey();
    }

    /**
     * 验证用户提交的 TOTP 码是否合法（基于用户的 Secret Key）
     *
     * @param secretKey Base32 编码的 TOTP 密钥
     * @param code      用户提交的 6 位验证码
     * @return true 如果校验通过；false 否则
     */
    public static boolean verifyTotp(String secretKey, int code) {
        return gAuth.authorize(secretKey, code);
    }

    /**
     * 将 Base32 编码的密钥转换为 Hex，若业务需要展示给前端 URI 可用该方法
     */
    public static String getHexKey(String base32Secret) {
        Base32 codec = new Base32();
        byte[] bytes = codec.decode(base32Secret);
        return Hex.encodeHexString(bytes);
    }

    /**
     * 生成在 Google Authenticator 中添加账户的二维码 URI
     *
     * @param username 用户名
     * @param secret   Base32 编码密钥
     * @param issuer   应用或企业名称，比如 "MyCompany"
     * @return otpauth://totp/issuer:username?secret=SECRET&issuer=issuer
     */
    public static String getOtpAuthURL(String username, String secret, String issuer) {
        return String.format(
                "otpauth://totp/%s:%s?secret=%s&issuer=%s",
                issuer, username, secret, issuer
        );
    }
}

public interface IUserService {
    User findByUsername(String username);
    void register(User user);
    void enableMfa(Long userId);
    boolean verifyTotp(Long userId, int code);
}

@Service
public class UserServiceImpl implements IUserService {

    @Autowired
    private UserMapper userMapper;

    private final BCryptPasswordEncoder passwordEncoder = new BCryptPasswordEncoder();

    @Override
    public User findByUsername(String username) {
        return userMapper.selectOne(new QueryWrapper<User>().eq("username", username));
    }

    @Override
    public void register(User user) {
        // 加密密码
        user.setPassword(passwordEncoder.encode(user.getPassword()));
        user.setEnabled(true);
        user.setMfaEnabled(false);
        user.setMfaSecret(null);
        userMapper.insert(user);
    }

    @Override
    public void enableMfa(Long userId) {
        // 为用户生成 TOTP Secret 并更新
        User u = userMapper.selectById(userId);
        String secret = TotpUtils.generateSecretKey();
        u.setMfaSecret(secret);
        u.setMfaEnabled(true);
        userMapper.updateById(u);
    }

    @Override
    public boolean verifyTotp(Long userId, int code) {
        User u = userMapper.selectById(userId);
        if (u == null || !u.getMfaEnabled() || u.getMfaSecret() == null) {
            return false;
        }
        return TotpUtils.verifyTotp(u.getMfaSecret(), code);
    }
}

public class MfaRequiredException extends AuthenticationException {
    private final String username;

    public MfaRequiredException(String msg, String username) {
        super(msg);
        this.username = username;
    }

    public String getUsername() {
        return username;
    }
}

/**
 * 第一步：校验用户名 + 密码
 * 如果用户启用 MFA，则抛出 MfaRequiredException，后续由 MfaAuthenticationFilter 处理
 */
@Component
public class CustomAuthenticationProvider implements AuthenticationProvider {

    @Autowired
    private IUserService userService;

    private final BCryptPasswordEncoder passwordEncoder = new BCryptPasswordEncoder();

    @Override
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        String username = authentication.getName();
        String password = (String) authentication.getCredentials();
        User user = userService.findByUsername(username);

        if (user == null || !user.getEnabled()) {
            throw new BadCredentialsException("用户名或密码错误");
        }
        if (!passwordEncoder.matches(password, user.getPassword())) {
            throw new BadCredentialsException("用户名或密码错误");
        }

        // 如果用户启用了 MFA，则抛出自定义异常，提示进行第二步验证
        if (Boolean.TRUE.equals(user.getMfaEnabled())) {
            throw new MfaRequiredException("MFA 验证必需", username);
        }

        // 未启用 MFA 则直接构建 Authentication
        return new UsernamePasswordAuthenticationToken(
                username, null,
                Collections.singletonList(new SimpleGrantedAuthority("ROLE_USER"))
        );
    }

    @Override
    public boolean supports(Class<?> authentication) {
        return UsernamePasswordAuthenticationToken.class.isAssignableFrom(authentication);
    }
}

/**
 * 该过滤器负责处理 /mfa-verify POST 请求，
 * 从 session 中获取待验证用户名，校验用户提交的 TOTP 码。
 */
@Component
public class MfaAuthenticationFilter extends AbstractAuthenticationProcessingFilter {

    private final IUserService userService;

    public MfaAuthenticationFilter(IUserService userService) {
        super(new AntPathRequestMatcher("/mfa-verify", "POST"));
        this.userService = userService;
        // 不让 Spring Security 为我们阻止 CSRF，示例中 CSRF 已关闭
    }

    @Override
    public Authentication attemptAuthentication(
            HttpServletRequest request,
            HttpServletResponse response
    ) throws IOException, ServletException {
        // 应用前端将用户名暂存到 sessionAttribute： "MFA_USER"
        String username = (String) request.getSession().getAttribute("MFA_USER");
        if (username == null) {
            throw new RuntimeException("会话中找不到待 MFA 用户");
        }

        // 获取用户提交的 TOTP 码
        String codeStr = request.getParameter("code");
        if (codeStr == null || codeStr.isEmpty()) {
            throw new RuntimeException("TOTP 码不能为空");
        }

        int code;
        try {
            code = Integer.parseInt(codeStr);
        } catch (NumberFormatException e) {
            throw new RuntimeException("TOTP 码格式不正确");
        }

        // 从数据库校验 TOTP
        User user = userService.findByUsername(username);
        boolean valid = userService.verifyTotp(user.getId(), code);
        if (!valid) {
            throw new RuntimeException("TOTP 验证失败");
        }

        // 验证成功，构建真正的 Authentication 对象
        UsernamePasswordAuthenticationToken auth =
                new UsernamePasswordAuthenticationToken(
                        username, null, Collections.singletonList(() -> "ROLE_USER")
                );
        return auth;
    }

    @Override
    protected void successfulAuthentication(
            HttpServletRequest request,
            HttpServletResponse response,
            FilterChain chain,
            Authentication authResult
    ) throws IOException, ServletException {
        // 将最终的 Authentication 填入 SecurityContext
        SecurityContextHolder.getContext().setAuthentication(authResult);
        // 登录成功后清除 session 中的 MFA 用户标志
        request.getSession().removeAttribute("MFA_USER");
        // 跳转到首页
        response.sendRedirect("/");
    }

    @Override
    protected void unsuccessfulAuthentication(
            HttpServletRequest request,
            HttpServletResponse response,
            AuthenticationException failed
    ) throws IOException, ServletException {
        // 验证失败，跳回 MFA 验证页面
        response.setContentType(MediaType.TEXT_PLAIN_VALUE);
        response.getWriter().write("MFA 验证失败：" + failed.getMessage());
        response.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
    }
}

/**
 * 处理第一步用户名/密码登录失败或触发 MFA 的情况
 */
@Component
public class CustomAuthenticationFailureHandler implements AuthenticationFailureHandler {

    @Override
    public void onAuthenticationFailure(
            HttpServletRequest request,
            HttpServletResponse response,
            AuthenticationException exception
    ) throws IOException, ServletException {
        // 如果是 MfaRequiredException，重定向到 /mfa 页面，并将用户名存入 session
        if (exception instanceof MfaRequiredException) {
            String username = ((MfaRequiredException) exception).getUsername();
            request.getSession().setAttribute("MFA_USER", username);
            // 重定向到 MFA 验证页面
            response.sendRedirect("/mfa");
        } else {
            // 普通登录失败，重定向回 /login?error
            response.sendRedirect("/login?error=true");
        }
    }
}

/**
 * 核心安全配置：
 *  1. 注入自定义 AuthenticationProvider
 *  2. 配置表单登录和 MfaAuthenticationFilter
 */
@Configuration
public class SecurityConfig {

    @Autowired
    private CustomAuthenticationProvider customAuthenticationProvider;

    @Autowired
    private MfaAuthenticationFilter mfaAuthenticationFilter;

    @Autowired
    private CustomAuthenticationFailureHandler customFailureHandler;

    @Bean
    public SecurityFilterChain filterChain(HttpSecurity http, AuthenticationConfiguration authConfig) throws Exception {
        // 禁用 CSRF 简化示例
        http.csrf(csrf -> csrf.disable());

        // 使用自定义 AuthenticationProvider 替换默认的 DaoAuthenticationProvider
        http.authenticationProvider(customAuthenticationProvider);

        // 1. 首先，配置表单登录
        http.authorizeHttpRequests(auth -> auth
                .requestMatchers("/login", "/register", "/css/**", "/js/**").permitAll()
                .anyRequest().authenticated()
        )
        .formLogin(form -> form
                .loginPage("/login")
                .loginProcessingUrl("/login") // 与表单提交 action 保持一致
                .failureHandler(customFailureHandler)
                .defaultSuccessUrl("/", true)
        );

        // 2. 注册 MFA 过滤器，它要在 UsernamePasswordAuthenticationFilter 之后执行
        http.addFilterAfter(mfaAuthenticationFilter, org.springframework.security.web.authentication.UsernamePasswordAuthenticationFilter.class);

        // 3. Session 管理：MFA 过程中会话保持
        http.sessionManagement(session -> session
                .sessionCreationPolicy(SessionCreationPolicy.IF_REQUIRED)
        );

        // 4. 未授权时返回 401
        http.exceptionHandling(ex -> ex.authenticationEntryPoint((request, response, authException) ->
                response.sendError(HttpServletResponse.SC_UNAUTHORIZED, "Unauthorized"))
        );

        return http.build();
    }

    // 若需手动获取 AuthenticationManager，可使用以下 Bean
    @Bean
    public AuthenticationManager authenticationManager(AuthenticationConfiguration authConfig) throws Exception {
        return authConfig.getAuthenticationManager();
    }
}

@Controller
public class AuthController {

    @Autowired
    private IUserService userService;

    /**
     * 登录页（第一步：用户名 + 密码）
     */
    @GetMapping("/login")
    public String loginPage(@RequestParam(required = false) String error, Model model) {
        model.addAttribute("error", error != null);
        return "login";
    }

    /**
     * 注册页（仅示例）
     */
    @GetMapping("/register")
    public String registerPage() {
        return "register";
    }

    @PostMapping("/register")
    public String doRegister(@RequestParam String username, @RequestParam String password) {
        User u = new User();
        u.setUsername(username);
        u.setPassword(password);
        userService.register(u);
        return "redirect:/login";
    }

    /**
     * MFA 验证页：用户输入动态验证码
     */
    @GetMapping("/mfa")
    public String mfaPage(HttpSession session, Model model) {
        String username = (String) session.getAttribute("MFA_USER");
        if (username == null) {
            // 无待验证用户，跳到登录页
            return "redirect:/login";
        }
        model.addAttribute("username", username);
        return "mfa";
    }

    /**
     * 启用 MFA：登录后用户请求此接口可获取 TOTP Secret 与二维码 URL
     */
    @GetMapping("/enable-mfa")
    public String enableMfa(Authentication authentication, Model model) {
        if (authentication == null || !authentication.isAuthenticated()) {
            return "redirect:/login";
        }
        String username = authentication.getName();
        User u = userService.findByUsername(username);
        if (u.getMfaEnabled()) {
            model.addAttribute("message", "MFA 已启用");
            return "home";
        }
        // 为用户生成秘钥并开启 MFA
        userService.enableMfa(u.getId());
        u = userService.findByUsername(username); // 刷新
        String secret = u.getMfaSecret();
        String otpAuthURL = TotpUtils.getOtpAuthURL(username, secret, "MyCompany");
        model.addAttribute("otpAuthURL", otpAuthURL);
        model.addAttribute("secret", secret);
        return "enable-mfa";
    }

    @GetMapping("/")
    public String homePage() {
        return "home";
    }
}