package cn.meowrain.aioj.backend.auth.oauth2.service;

import cn.meowrain.aioj.backend.auth.oauth2.exception.InvalidGrantException;
import cn.meowrain.aioj.backend.auth.oauth2.exception.OAuth2Exception;
import cn.hutool.core.util.StrUtil;
import org.springframework.stereotype.Service;

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

/**
 * PKCE (Proof Key for Code Exchange) 验证服务 实现 RFC 7636 规范
 *
 * @author meowrain
 * @since 2025-12-14
 */
@Service
public class PKCEService {

	private static final String S256_METHOD = "S256";

	private static final String PLAIN_METHOD = "plain";

	/**
	 * 验证 PKCE
	 * @param codeVerifier Code Verifier（客户端生成的随机字符串，43-128字符）
	 * @param codeChallenge Code Challenge（授权请求时发送的挑战值）
	 * @param codeChallengeMethod 挑战方法（S256 或 plain）
	 * @throws InvalidGrantException 验证失败时抛出
	 */
	public void validatePKCE(String codeVerifier, String codeChallenge, String codeChallengeMethod) {
		// 1. 验证参数
		if (StrUtil.isBlank(codeVerifier)) {
			throw new InvalidGrantException("code_verifier 不能为空");
		}

		if (StrUtil.isBlank(codeChallenge)) {
			throw new InvalidGrantException("code_challenge 不能为空");
		}

		if (StrUtil.isBlank(codeChallengeMethod)) {
			throw new InvalidGrantException("code_challenge_method 不能为空");
		}

		// 2. 验证 code_verifier 长度（RFC 7636: 43-128 字符）
		if (codeVerifier.length() < 43 || codeVerifier.length() > 128) {
			throw new InvalidGrantException("code_verifier 长度必须在 43-128 字符之间");
		}

		// 3. 根据方法验证
		String computedChallenge;
		switch (codeChallengeMethod) {
			case S256_METHOD:
				computedChallenge = computeS256Challenge(codeVerifier);
				break;
			case PLAIN_METHOD:
				// plain 方法：challenge = verifier（不推荐使用）
				computedChallenge = codeVerifier;
				break;
			default:
				throw new OAuth2Exception("不支持的 code_challenge_method: " + codeChallengeMethod);
		}

		// 4. 比对 challenge
		if (!computedChallenge.equals(codeChallenge)) {
			throw new InvalidGrantException("code_verifier 验证失败");
		}
	}

	/**
	 * 计算 S256 方法的 Code Challenge challenge = BASE64URL(SHA256(verifier))
	 * @param codeVerifier Code Verifier
	 * @return Code Challenge
	 */
	private String computeS256Challenge(String codeVerifier) {
		try {
			MessageDigest digest = MessageDigest.getInstance("SHA-256");
			byte[] hash = digest.digest(codeVerifier.getBytes(StandardCharsets.US_ASCII));
			return Base64.getUrlEncoder().withoutPadding().encodeToString(hash);
		}
		catch (NoSuchAlgorithmException e) {
			throw new RuntimeException("SHA-256 算法不可用", e);
		}
	}

	/**
	 * 验证 PKCE 参数（在授权请求阶段）
	 * @param codeChallenge Code Challenge
	 * @param codeChallengeMethod 挑战方法
	 * @param requirePkce 是否强制要求 PKCE
	 * @throws OAuth2Exception 验证失败时抛出
	 */
	public void validatePKCERequest(String codeChallenge, String codeChallengeMethod, boolean requirePkce) {
		// 如果强制要求 PKCE
		if (requirePkce) {
			if (StrUtil.isBlank(codeChallenge)) {
				throw new OAuth2Exception("该客户端必须使用 PKCE，请提供 code_challenge");
			}
			if (StrUtil.isBlank(codeChallengeMethod)) {
				throw new OAuth2Exception("该客户端必须使用 PKCE，请提供 code_challenge_method");
			}
		}

		// 如果提供了 code_challenge，验证 method
		if (StrUtil.isNotBlank(codeChallenge)) {
			if (StrUtil.isBlank(codeChallengeMethod)) {
				throw new OAuth2Exception("提供了 code_challenge 必须同时提供 code_challenge_method");
			}

			// 只支持 S256 方法（更安全）
			if (!S256_METHOD.equals(codeChallengeMethod) && !PLAIN_METHOD.equals(codeChallengeMethod)) {
				throw new OAuth2Exception("不支持的 code_challenge_method，仅支持 S256 或 plain");
			}

			// 推荐使用 S256
			if (PLAIN_METHOD.equals(codeChallengeMethod)) {
				// 可以记录警告日志：使用了不安全的 plain 方法
			}
		}
	}

}
