package com.benu.anubis.salesforce.service.impl;

import com.benu.anubis.core.exception.ExceptionReporterService;
import com.benu.anubis.salesforce.service.api.HttpRequestHandler;
import com.benu.anubis.salesforce.service.exception.AuthenticationException;
import com.benu.anubis.salesforce.service.api.OAuthService;
import com.benu.anubis.salesforce.entity.OauthResponse;
import com.benu.anubis.salesforce.repository.AccessTokenRepository;
import org.apache.commons.codec.binary.Base64;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.FileInputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.Signature;
import java.text.MessageFormat;
import java.time.LocalDateTime;
import java.util.Optional;

import static org.springframework.http.HttpMethod.POST;

/**
 * SalesForce uses the Oauth strategy to implement authentication. First you will need to register the client,
 * then you will need to create a self-client token. With that self-client token you can make a first token generation
 * request to get a refresh_token. With that refresh_token, you can then generate access_tokens as you make
 * requests to the platform. All these steps were done and the refresh_token is saved as a secret and inject in the
 * SalesForce.token.refresh_token property in the catalog-service.yml configuration file for each environment.
 * Check out the official documentation here: https://www.SalesForce.com/crm/help/developer/api/oauth-overview.html
 */
@Service
public class OAuthServiceImpl implements OAuthService {

  private final Logger logger = LoggerFactory.getLogger(OAuthServiceImpl.class);


  String header = "{\"alg\":\"RS256\"}";
  String claimTemplate = "'{'\"iss\": \"{0}\", \"sub\": \"{1}\", \"aud\": \"{2}\", \"exp\": \"{3}\"'}'";

  @Value("${oauth2.jwt.iss}")
  private String iss;
  @Value("${oauth2.jwt.sub}")
  private String sub;
  @Value("${oauth2.jwt.aud}")
  private String aud;
  @Value("${oauth2.jwt.jks.location}")
  private String jksLocation;
  @Value("${oauth2.jwt.jks.key}")
  private String jksKey;
  @Value("${oauth2.jwt.jks.alias}")
  private String jksAlias;
  @Value("${oauth2.token}")
  private String tokenUrl;
  @Value("${oauth2.grantType}")
  private String grantType;

  private final RestTemplate restTemplate;
  private final AccessTokenRepository accessTokenRepository;
  private final ExceptionReporterService exceptionReporterService;

  @Autowired
  public OAuthServiceImpl(RestTemplate restTemplate, AccessTokenRepository accessTokenRepository, ExceptionReporterService exceptionReporterService) {
    this.restTemplate = restTemplate;
    this.accessTokenRepository = accessTokenRepository;
    this.exceptionReporterService = exceptionReporterService;
  }

  /**
   * Creates the HTTP header with the Authorization entry (cached access token/new access token)
   * @return http headers with the authorization attribute
   * @throws AuthenticationException In case there was an issue generating the access token
   */
  @NotNull
  public HttpHeaders getAuthHeader() throws AuthenticationException {
    HttpHeaders headers = new HttpHeaders();
    headers.add(HttpHeaders.AUTHORIZATION, "Bearer " + getAccessToken());
    return headers;
  }

  /**
   * Creates the HTTP header with the Authorization entry (new access token)
   * @return http headers with the authorization attribute
   * @throws AuthenticationException In case there was an issue generating the access token
   */
  @NotNull
  public HttpHeaders getAuthHeaderWithNewToken() throws AuthenticationException {
    HttpHeaders headers = new HttpHeaders();
    headers.add(HttpHeaders.AUTHORIZATION, "Bearer " + generateNewAccessToken().getAccessToken());
    return headers;
  }

  private String getAccessToken() throws AuthenticationException {
    try {
      Optional<OauthResponse> tokenFromDatabase = findAvailableAccessToken();
      return tokenFromDatabase.isPresent() ?
              tokenFromDatabase.get().getAccessToken()
              : generateNewAccessToken().getAccessToken();
    } catch (Exception exception) {
      logger.error(String.format("[Salesforce OAuth Service] Failed to Authenticate: %s", exception.getMessage()));
      exceptionReporterService.report(exception, "generateAccessToken");
      throw new AuthenticationException(exception.getMessage());
    }
  }

  private Optional<OauthResponse> findAvailableAccessToken() {
    return accessTokenRepository.findTopByExpirationDateGreaterThan(LocalDateTime.now());
  }

  /**
   * Generate the token and save it to be reused in the future
   * @return  OOauthResponse the generated token
   */
  private OauthResponse generateNewAccessToken() throws AuthenticationException {
    UriComponentsBuilder builder =
            UriComponentsBuilder.fromHttpUrl(tokenUrl)
                    .queryParam("grant_type", grantType)
                    .queryParam("assertion", generateJWT());

    ParameterizedTypeReference<OauthResponse> ptr = new ParameterizedTypeReference<>() {};
    ResponseEntity<OauthResponse> response = restTemplate.exchange(builder.toUriString(), POST, null, ptr);
    OauthResponse responseBody = Optional.ofNullable(response.getBody()).orElseThrow(RuntimeException::new);
    saveAccessToken(responseBody);
    return responseBody;
  }

  /**
   * Saves the token to the database, setting the expiration date for 2 hour later so that it can be easily retrieved
   * and used up until it expires.
   * @param accessToken the generated token
   */
  private void saveAccessToken(OauthResponse accessToken) {
    LocalDateTime localDateTime = LocalDateTime.now().plusHours(2);
    accessToken.setExpirationDate(localDateTime);
    accessTokenRepository.save(accessToken);
  }

  /**
   * Generate JWT token
   * Check out documentation here: https://help.salesforce.com/articleView?id=sf.remoteaccess_oauth_jwt_flow.htm&type=5
   * @return JWT generated token
   */
  public String generateJWT() throws AuthenticationException {
    try {
      StringBuilder token = new StringBuilder();

      //Encode the JWT Header and add it to our string to sign
      token.append(Base64.encodeBase64URLSafeString(header.getBytes(StandardCharsets.UTF_8)));

      //Separate with a period
      token.append(".");

      //Create the JWT Claims Object
      String[] claimArray = new String[4];
      claimArray[0] = iss;
      claimArray[1] = sub;
      claimArray[2] = aud;
      claimArray[3] = Long.toString((System.currentTimeMillis() / 1000) + 300);

      MessageFormat claims;
      claims = new MessageFormat(claimTemplate);
      String payload = claims.format(claimArray);

      //Add the encoded claims object
      token.append(Base64.encodeBase64URLSafeString(payload.getBytes(StandardCharsets.UTF_8)));

      //Load the private key from a keystore
      KeyStore keystore = KeyStore.getInstance("JKS");
      keystore.load(new FileInputStream(jksLocation), jksKey.toCharArray());
      PrivateKey privateKey = (PrivateKey) keystore.getKey(jksAlias, jksKey.toCharArray());

      //Sign the JWT Header + "." + JWT Claims Object
      Signature signature = Signature.getInstance("SHA256withRSA");
      signature.initSign(privateKey);
      signature.update(token.toString().getBytes(StandardCharsets.UTF_8));
      String signedPayload = Base64.encodeBase64URLSafeString(signature.sign());

      //Separate with a period
      token.append(".");

      //Add the encoded signature
      token.append(signedPayload);
      return token.toString();
    } catch (Exception e) {
      throw new AuthenticationException("Error While Generating JWT Token");
    }
  }
}