/*
 * Copyright 2025 Hirokazu Kobayashi
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.idp.server.core.openid.extension.pkce;

import org.idp.server.core.openid.grant_management.grant.AuthorizationCodeGrant;
import org.idp.server.core.openid.oauth.clientauthenticator.clientcredentials.ClientCredentials;
import org.idp.server.core.openid.oauth.request.AuthorizationRequest;
import org.idp.server.core.openid.oauth.type.pkce.CodeChallenge;
import org.idp.server.core.openid.oauth.type.pkce.CodeVerifier;
import org.idp.server.core.openid.token.TokenRequestContext;
import org.idp.server.core.openid.token.exception.TokenBadRequestException;
import org.idp.server.core.openid.token.verifier.AuthorizationCodeGrantExtensionVerifierInterface;
import org.idp.server.platform.log.LoggerWrapper;

public class AuthorizationCodeGrantPkceVerifier
    implements AuthorizationCodeGrantExtensionVerifierInterface {

  LoggerWrapper log = LoggerWrapper.getLogger(AuthorizationCodeGrantPkceVerifier.class);

  @Override
  public boolean shouldVerify(
      TokenRequestContext tokenRequestContext,
      AuthorizationRequest authorizationRequest,
      AuthorizationCodeGrant authorizationCodeGrant,
      ClientCredentials clientCredentials) {
    // Issue #1523: also enforce when the client mandates PKCE, so a missing code_verifier is
    // rejected with invalid_grant even if the (defensively) stored request lacked a code_challenge.
    return authorizationRequest.isPkceRequest()
        || tokenRequestContext.clientConfiguration().isRequirePkce();
  }

  @Override
  public void verify(
      TokenRequestContext tokenRequestContext,
      AuthorizationRequest authorizationRequest,
      AuthorizationCodeGrant authorizationCodeGrant,
      ClientCredentials clientCredentials) {

    log.debug("AuthorizationCodeGrantPkceVerifier verification start");
    throwExceptionIfNotContainsCodeVerifier(tokenRequestContext);
    throwExceptionIfUnMatchCodeVerifier(tokenRequestContext, authorizationRequest);
    throwExceptionIfInvalidCodeVerifierFormat(tokenRequestContext);
    log.debug("AuthorizationCodeGrantPkceVerifier verification end");
  }

  /**
   * RFC 7636 Section 4.6:
   *
   * <p>If the server requires Proof Key for Code Exchange (PKCE) and the client does not send the
   * "code_verifier" in the Token Request, the authorization server MUST return an "invalid_grant"
   * error.
   */
  void throwExceptionIfNotContainsCodeVerifier(TokenRequestContext tokenRequestContext) {
    if (!tokenRequestContext.hasCodeVerifier()) {
      throw new TokenBadRequestException(
          "invalid_grant",
          "authorization request has code_challenge, but token request does not contains code verifier");
    }
  }

  /**
   * RFC 7636 Section 4.6:
   *
   * <p>If the values are not equal, return an "invalid_grant" error response.
   */
  void throwExceptionIfUnMatchCodeVerifier(
      TokenRequestContext tokenRequestContext, AuthorizationRequest authorizationRequest) {
    if (authorizationRequest.isPkceWithS256()) {
      CodeVerifier codeVerifier = tokenRequestContext.codeVerifier();
      CodeChallengeCalculator codeChallengeCalculator = new CodeChallengeCalculator(codeVerifier);
      CodeChallenge codeChallenge = codeChallengeCalculator.calculateWithS256();
      if (!codeChallenge.equals(authorizationRequest.codeChallenge())) {
        throw new TokenBadRequestException(
            "invalid_grant",
            "code_verifier of token request does not match code_challenge of authorization request");
      }
      return;
    }
    CodeChallengeCalculator codeChallengeCalculator =
        new CodeChallengeCalculator(tokenRequestContext.codeVerifier());
    CodeChallenge codeChallenge = codeChallengeCalculator.calculateWithPlain();
    if (!codeChallenge.equals(authorizationRequest.codeChallenge())) {
      throw new TokenBadRequestException(
          "invalid_grant",
          "code_verifier of token request does not match code_challenge of authorization request");
    }
  }

  void throwExceptionIfInvalidCodeVerifierFormat(TokenRequestContext tokenRequestContext) {
    CodeVerifier codeVerifier = tokenRequestContext.codeVerifier();
    if (codeVerifier.isShorterThan43()) {
      throw new TokenBadRequestException(
          "invalid_grant", "code_verifier must be at least 43 characters");
    }

    if (codeVerifier.isLongerThan128()) {
      throw new TokenBadRequestException(
          "invalid_grant", "code_verifier must be at most 128 characters");
    }
    if (!codeVerifier.value().matches("^[A-Za-z0-9\\-._~]+$")) {
      throw new TokenBadRequestException(
          "invalid_grant", "code_verifier contains invalid characters");
    }
  }
}