package org.ta.common.util;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * PKCE 工具类（符合 RFC 7636 规范，适配 OAuth2 1.5.2）
 * 核心功能：生成 code_verifier、计算 code_challenge、验证 code_verifier 与 code_challenge 一致性
 */
public class PkceUtil {

    // code_verifier 最小长度（RFC 7636 要求：43-128 位）
    private static final int MIN_VERIFIER_LENGTH = 43;
    private static final int MAX_VERIFIER_LENGTH = 128;
    // 默认生成 64 位的 code_verifier（推荐长度）
    private static final int DEFAULT_VERIFIER_LENGTH = 64;
    // 随机字符串的字符集（RFC 7636 允许：a-z, A-Z, 0-9, -, ., _, ~）
    private static final String VERIFIER_CHARSET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~";
    private static final SecureRandom SECURE_RANDOM = new SecureRandom();

    /**
     * 生成合规的 code_verifier（随机字符串，长度 43-128 位）
     * @return code_verifier
     */
    public static String generateCodeVerifier() {
        return generateCodeVerifier(DEFAULT_VERIFIER_LENGTH);
    }

    /**
     * 生成指定长度的 code_verifier
     * @param length 长度（必须 43 ≤ length ≤ 128）
     * @return code_verifier
     * @throws IllegalArgumentException 长度不符合要求时抛出
     */
    public static String generateCodeVerifier(int length) {
        if (length < MIN_VERIFIER_LENGTH || length > MAX_VERIFIER_LENGTH) {
            throw new IllegalArgumentException("Code verifier length must be between 43 and 128");
        }

        StringBuilder verifier = new StringBuilder(length);
        int charsetLength = VERIFIER_CHARSET.length();
        for (int i = 0; i < length; i++) {
            verifier.append(VERIFIER_CHARSET.charAt(SECURE_RANDOM.nextInt(charsetLength)));
        }
        return verifier.toString();
    }

    /**
     * 根据 code_verifier 计算 code_challenge（SHA256 + Base64URL 编码，无填充）
     * @param codeVerifier code_verifier
     * @return code_challenge
     * @throws NoSuchAlgorithmException SHA256 算法不可用时抛出
     */
    public static String generateCodeChallenge(String codeVerifier) throws NoSuchAlgorithmException {
        // 1. 对 code_verifier 做 SHA256 哈希
        MessageDigest digest = MessageDigest.getInstance("SHA-256");
        byte[] hash = digest.digest(codeVerifier.getBytes(StandardCharsets.US_ASCII));

        // 2. Base64URL 编码（替换 +→-、/→_，移除末尾 = 填充）
        return Base64.getUrlEncoder()
                .withoutPadding() // 移除填充（RFC 7636 要求）
                .encodeToString(hash);
    }

    /**
     * 验证 code_verifier 是否与 code_challenge 匹配（用于授权服务器端验证）
     * @param codeVerifier 客户端提交的 code_verifier
     * @param codeChallenge 授权请求中的 code_challenge
     * @param challengeMethod 挑战方法（仅支持 S256，1.5.2 强制要求）
     * @return true=匹配，false=不匹配
     * @throws NoSuchAlgorithmException SHA256 算法不可用时抛出
     */
    public static boolean verifyCodeVerifier(String codeVerifier, String codeChallenge, String challengeMethod) throws NoSuchAlgorithmException {
        // 1.5.2 强制要求 S256 方法，不支持 plain
        if (!"S256".equalsIgnoreCase(challengeMethod)) {
            return false;
        }

        // 计算 code_verifier 对应的 code_challenge，与传入的对比
        String computedChallenge = generateCodeChallenge(codeVerifier);
        return computedChallenge.equals(codeChallenge);
    }

    // 测试主方法（验证工具类功能）
    public static void main(String[] args) throws NoSuchAlgorithmException {
        // 1. 生成 code_verifier
        String codeVerifier = generateCodeVerifier();
        System.out.println("1. 生成的 Code Verifier：" + codeVerifier);
        System.out.println("   长度：" + codeVerifier.length());

        // 2. 计算 code_challenge
        String codeChallenge = generateCodeChallenge(codeVerifier);
        System.out.println("\n2. 计算的 Code Challenge：" + codeChallenge);

        // 3. 验证匹配性
        boolean isMatch = verifyCodeVerifier(codeVerifier, codeChallenge, "S256");
        System.out.println("\n3. 验证结果（匹配）：" + isMatch); // 应输出 true

        // 4. 验证不匹配场景
        boolean isNotMatch = verifyCodeVerifier("wrong_verifier_1234567890", codeChallenge, "S256");
        System.out.println("4. 验证结果（不匹配）：" + isNotMatch); // 应输出 false
    }
}