package org.psesquared.server.config;

import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.MalformedJwtException;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.UnsupportedJwtException;
import io.jsonwebtoken.io.Decoders;
import io.jsonwebtoken.security.Keys;
import io.jsonwebtoken.security.SignatureException;
import java.security.Key;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;
import lombok.RequiredArgsConstructor;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.stereotype.Service;

/**
 * The service class responsible for creating, managing and validating JWTs.
 */
@Service
@RequiredArgsConstructor
public class JwtService {

  /**
   * The boolean value for expressing that a JWT is not valid.
   */
  private static final boolean INVALID = false;

  /**
   * The 1h lifespan of an access token.
   */
  private static final long ACCESS_TOKEN_LIFESPAN_MILLIS
      = 1000 * 60 * (long) 60;

  /**
   * The 24h lifespan of a URL token (verification/resetting password).
   */
  private static final long URL_TOKEN_LIFESPAN_MILLIS
      = ACCESS_TOKEN_LIFESPAN_MILLIS * 24;

  /**
   * The properties class that is used to return externally stored signing key.
   */
  private final SecurityConfigProperties securityConfigProperties;

  /**
   * Extracts the username from a JWT for authentication.
   *
   * @param token The JWT
   * @return The extracted username
   */
  public String extractAuthUsername(final String token) {
    return extractClaim(token, getAuthSigningKey(), Claims::getSubject);
  }

  /**
   * Extracts a generic claim from the JWT.
   *
   * @param <T> The type of the claim
   * @param token The JWT
   * @param signingKey The JWT signing key
   * @param claimsResolver The function to resolve the claim
   * @return The extracted generic claim
   * @throws ExpiredJwtException If the JWT has expired
   * @throws UnsupportedJwtException If the JWT is not supported
   * @throws MalformedJwtException If the JWT is malformed
   * @throws SignatureException If the signature doesn't match
   * @throws IllegalArgumentException If the token has an inappropriate format
   */
  public <T> T extractClaim(final String token,
                            final Key signingKey,
                            final Function<Claims, T> claimsResolver)
      throws ExpiredJwtException,
      UnsupportedJwtException,
      MalformedJwtException,
      SignatureException,
      IllegalArgumentException {

    final Claims claims = extractAllClaims(token, signingKey);
    return claimsResolver.apply(claims);
  }

  /**
   * Generates the JWT with additional claims and a lifespan for the
   * {@link org.psesquared.server.model.User} with the given details.
   *
   * @param additionalClaims The {@link Map} with additional claims
   * @param userDetails The user details
   * @param tokenLifespan The lifespan of the token
   * @param signingKey The JWT signing key
   * @return The generated JWT
   */
  public String generateTokenString(final Map<String, Object> additionalClaims,
                                    final UserDetails userDetails,
                                    final long tokenLifespan,
                                    final Key signingKey) {
    return Jwts.builder()
        .setClaims(additionalClaims)
        .setSubject(userDetails.getUsername())
        .setIssuedAt(new Date())
        .setExpiration(new Date(System.currentTimeMillis() + tokenLifespan))
        .signWith(signingKey, SignatureAlgorithm.HS256)
        .compact();
  }

  /**
   * Generates a JWT access token for the
   * {@link org.psesquared.server.model.User} with the given details.
   * (no additional claims).
   *
   * @param userDetails The user details
   * @return The generated JWT access token
   */
  public String generateAccessTokenString(final UserDetails userDetails) {
    //no additional claims supported but open for extension with roles e.g.
    return generateTokenString(new HashMap<>(),
        userDetails,
        ACCESS_TOKEN_LIFESPAN_MILLIS,
        getAuthSigningKey());
  }

  /**
   * Generates a JWT URL token required for authentication/resetting password
   * for the {@link org.psesquared.server.model.User} with the given details
   * (no additional claims).
   *
   * @param userDetails The user details
   * @return The generated JWT access token
   */
  public String generateUrlTokenString(final UserDetails userDetails) {
    return generateTokenString(new HashMap<>(),
        userDetails,
        URL_TOKEN_LIFESPAN_MILLIS,
        getUrlSigningKey());
  }

  /**
   * Validates the given JWT for authentication against the given
   * {@link UserDetails} and checks if it has not expired.
   *
   * @param token The to be validated JWT
   * @param userDetails The user details
   * @return {@code true} if the JWT is valid,
   *        <br>
   *        {@code false} otherwise
   */
  public boolean isAuthTokenValid(final String token,
                                  final UserDetails userDetails) {
    return isTokenValid(token, userDetails, getAuthSigningKey());
  }

  /**
   * Validates the given JWT for URLs against the given {@link UserDetails}
   * and checks if it has not expired.
   *
   * @param token The to be validated JWT
   * @param userDetails The user details
   * @return {@code true} if the JWT is valid,
   *        <br>
   *        {@code false} otherwise
   */
  public boolean isUrlTokenValid(final String token,
                                  final UserDetails userDetails) {
    return isTokenValid(token, userDetails, getUrlSigningKey());
  }

  /**
   * Validates the given JWT against the given {@link UserDetails}
   * with the given signing key and checks if it has not expired.
   *
   * @param token The to be validated JWT
   * @param userDetails The user details
   * @param signingKey The JWT signing key
   * @return {@code true} if the JWT is valid,
   *        <br>
   *        {@code false} otherwise
   */
  private boolean isTokenValid(final String token,
                               final UserDetails userDetails,
                               final Key signingKey) {
    try {
      final String username = extractUsername(token, signingKey);
      return username.equals(userDetails.getUsername())
          && !isTokenExpired(token, signingKey);
    } catch (ExpiredJwtException
             | UnsupportedJwtException
             | MalformedJwtException
             | SignatureException
             | IllegalArgumentException e) {
      return INVALID;
    }
  }

  /**
   * Checks if the given JWT is expired.
   *
   * @param token The JWT
   * @param signingKey The JWT signing key
   * @return {@code true} if the JWT is expired,
   *        <br>
   *        {@code false} otherwise
   */
  private boolean isTokenExpired(final String token, final Key signingKey) {
    return extractExpiration(token, signingKey).before(new Date());
  }

  /**
   * Extracts the username from a JWT.
   *
   * @param token The JWT
   * @param signingKey The JWT signing key
   * @return The extracted username
   */
  private String extractUsername(final String token, final Key signingKey) {
    return extractClaim(token, signingKey, Claims::getSubject);
  }

  /**
   * Extracts the expiration {@link Date} of the JWT.
   *
   * @param token The JWT
   * @param signingKey The JWT signing key
   * @return The expiration date
   */
  private Date extractExpiration(final String token, final Key signingKey) {
    return extractClaim(token, signingKey, Claims::getExpiration);
  }

  /**
   * Extracts all claims from the JWT in order for
   * {@link #extractClaim(String, Key, Function)} to be able
   * to filter out one claim.
   *
   * @param token The JWT
   * @param signingKey The JWT signing key
   * @return All claims of the JWT
   * @throws ExpiredJwtException If the JWT has expired
   * @throws UnsupportedJwtException If the JWT is not supported
   * @throws MalformedJwtException If the JWT is malformed
   * @throws SignatureException If the signature doesn't match
   * @throws IllegalArgumentException If the token has an inappropriate format
   */
  private Claims extractAllClaims(final String token, final Key signingKey)
      throws ExpiredJwtException,
      UnsupportedJwtException,
      MalformedJwtException,
      SignatureException,
      IllegalArgumentException {

    return Jwts.parserBuilder()
        .setSigningKey(signingKey)
        .build()
        .parseClaimsJws(token)
        .getBody();
  }

  /**
   * Returns the signing {@link Key} for signing the JWT.
   *
   * @return The signing key
   */
  private Key getAuthSigningKey() {
    byte[] keyBytes
        = Decoders.BASE64.decode(securityConfigProperties.jwtAuthSigningKey());
    return Keys.hmacShaKeyFor(keyBytes);
  }

  /**
   * Returns the signing {@link Key} for signing the JWT.
   *
   * @return The signing key
   */
  private Key getUrlSigningKey() {
    byte[] keyBytes
        = Decoders.BASE64.decode(securityConfigProperties.jwtUrlSigningKey());
    return Keys.hmacShaKeyFor(keyBytes);
  }

}