package com.eactive.eai.authserver.jwt;

import java.lang.reflect.Field;
import java.nio.charset.StandardCharsets;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Base64;
import java.util.LinkedHashMap;
import java.util.Map;

import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;

import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * PS256(RSA-PSS) JWT를 발급·검증하는 JwtAccessTokenConverter 확장 클래스.
 *
 * spring-security-jwt 1.1.1은 JwtAlgorithms에 PS256을 등록하지 않아
 * JwtHelper.encode/decodeAndVerify가 PS256 헤더를 처리하지 못한다.
 * 따라서 encode/decode를 직접 구현하여 JwtHelper를 완전히 우회한다.
 */
public class PssJwtAccessTokenConverter extends JwtAccessTokenConverter {

    /** PS256 JWT 헤더 {"alg":"PS256","typ":"JWT"} 의 Base64URL 인코딩 (고정값) */
    private static final String ENCODED_HEADER = Base64.getUrlEncoder().withoutPadding()
            .encodeToString("{\"alg\":\"PS256\",\"typ\":\"JWT\"}".getBytes(StandardCharsets.UTF_8));

    private final RsaPssSigner   pssSigner;
    private final RsaPssVerifier pssVerifier;
    private final String         publicKeyPem;
    private final ObjectMapper   jackson = new ObjectMapper();

    public PssJwtAccessTokenConverter(RSAPrivateKey privateKey, RSAPublicKey publicKey) {
        this.pssSigner    = new RsaPssSigner(privateKey);
        this.pssVerifier  = new RsaPssVerifier(publicKey);
        this.publicKeyPem = "-----BEGIN PUBLIC KEY-----\n"
                + Base64.getMimeEncoder(64, new byte[]{'\n'}).encodeToString(publicKey.getEncoded())
                + "\n-----END PUBLIC KEY-----";
        // 부모 setVerifierKey — afterPropertiesSet() 내부에서 참조하는 verifierKey 필드 초기화
        setVerifierKey(this.publicKeyPem);
        // getKey() (/oauth/token_key 엔드포인트)가 동작하도록 부모의 verifier 필드를 주입
        try {
            Field f = JwtAccessTokenConverter.class.getDeclaredField("verifier");
            f.setAccessible(true);
            f.set(this, pssVerifier);
        } catch (Exception ignored) { }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        // 부모 afterPropertiesSet()은 signingKey 기반 초기화(JwtHelper 의존)를 시도하므로 건너뜀
    }

    /**
     * PS256 JWT 발급. JwtHelper.encode를 사용하지 않고 직접 구성한다.
     * signingInput = Base64URL(header) + "." + Base64URL(payload)
     * JWT = signingInput + "." + Base64URL(PSS_sign(signingInput))
     */
    @Override
    protected String encode(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        Map<String, ?> claims = getAccessTokenConverter().convertAccessToken(accessToken, authentication);
        String payload;
        try {
            payload = jackson.writeValueAsString(claims);
        } catch (Exception e) {
            throw new IllegalStateException("Access token을 JSON으로 변환 실패", e);
        }
        return buildJwt(payload);
    }

    /**
     * PS256 JWT 검증 및 클레임 추출. JwtHelper.decodeAndVerify를 사용하지 않고 직접 파싱한다.
     */
    @Override
    protected Map<String, Object> decode(String token) {
        String[] parts = token.split("\\.");
        if (parts.length != 3) {
            throw new IllegalArgumentException("잘못된 JWT 형식");
        }
        byte[] signingInput = (parts[0] + "." + parts[1]).getBytes(StandardCharsets.US_ASCII);
        byte[] sig          = Base64.getUrlDecoder().decode(padBase64Url(parts[2]));
        pssVerifier.verify(signingInput, sig); // 실패 시 InvalidSignatureException

        byte[] payloadBytes = Base64.getUrlDecoder().decode(padBase64Url(parts[1]));
        try {
            @SuppressWarnings("unchecked")
            Map<String, Object> result = jackson.readValue(payloadBytes, Map.class);
            // JwtAccessTokenConverter 원본 동작 유지: exp 필드 Integer → Long 변환
            if (result.containsKey(EXP) && result.get(EXP) instanceof Integer) {
                result.put(EXP, ((Integer) result.get(EXP)).longValue());
            }
            return result;
        } catch (Exception e) {
            throw new IllegalArgumentException("JWT 페이로드 파싱 실패", e);
        }
    }

    /** /oauth/token_key 엔드포인트 응답 — alg와 PEM 공개키 반환 */
    @Override
    public Map<String, String> getKey() {
        Map<String, String> result = new LinkedHashMap<>();
        result.put("alg",   "PS256");
        result.put("value", publicKeyPem);
        return result;
    }

    public String buildJwt(String payload) {
        String encodedPayload = Base64.getUrlEncoder().withoutPadding()
                .encodeToString(payload.getBytes(StandardCharsets.UTF_8));
        String signingInput = ENCODED_HEADER + "." + encodedPayload;
        byte[] sig = pssSigner.sign(signingInput.getBytes(StandardCharsets.US_ASCII));
        return signingInput + "." + Base64.getUrlEncoder().withoutPadding().encodeToString(sig);
    }

    private static String padBase64Url(String s) {
        int mod = s.length() % 4;
        if (mod == 2) return s + "==";
        if (mod == 3) return s + "=";
        return s;
    }
}
