package com.airlock.iam.flow.shared.infrastructure.service.newdevice;

import com.airlock.iam.base.api.domain.model.exception.FailureException;
import com.airlock.iam.base.api.domain.model.exception.ViolationException;
import com.airlock.iam.base.api.domain.model.user.UserId;
import com.airlock.iam.base.misc.collection.Pair;
import com.airlock.iam.base.misc.log.LogUtil;
import com.airlock.iam.base.misc.result.Result;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceCreatedAt;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceEncryptionKey;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceKnownUsers;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceRenewedAt;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceSettings;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceUser;
import com.airlock.iam.flow.shared.domain.model.newdevice.LoginFromNewDeviceUserInfo;
import com.nimbusds.jose.EncryptionMethod;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JOSEObjectType;
import com.nimbusds.jose.JWEAlgorithm;
import com.nimbusds.jose.JWEHeader;
import com.nimbusds.jose.crypto.DirectDecrypter;
import com.nimbusds.jose.crypto.DirectEncrypter;
import com.nimbusds.jose.util.JSONObjectUtils;
import com.nimbusds.jwt.EncryptedJWT;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.JWTClaimsSet.Builder;
import com.nimbusds.jwt.util.DateUtils;
import jakarta.inject.Inject;
import jakarta.inject.Provider;
import java.text.ParseException;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;

public class LoginFromNewDeviceKnownUsersSerializer {
   static final JWEAlgorithm DEFAULT_ALGORITHM = JWEAlgorithm.DIR;
   static final EncryptionMethod ENCRYPTION_METHOD = EncryptionMethod.A256GCM;
   static final String JWT_CLAIM_USERS = "users";
   static final String JSON_KEY_CREATED_AT = "created";
   static final String JSON_KEY_RENEWED_AT = "renewed";
   private final Provider<LoginFromNewDeviceSettings> settings;

   @Inject
   public LoginFromNewDeviceKnownUsersSerializer(Provider<LoginFromNewDeviceSettings> settings) {
      this.settings = settings;
   }

   public String serialize(LoginFromNewDeviceKnownUsers knownUsers) {
      Map<String, Object> users = knownUsers.users()
         .entrySet()
         .stream()
         .map(LoginFromNewDeviceKnownUsersSerializer::toJSONSerializable)
         .collect(Collectors.toMap(Pair::getFirst, Pair::getSecond));
      JWTClaimsSet claimsSet = new Builder().claim("users", users).build();
      JWEHeader header = new com.nimbusds.jose.JWEHeader.Builder(DEFAULT_ALGORITHM, ENCRYPTION_METHOD).type(JOSEObjectType.JWT).build();
      EncryptedJWT jwt = new EncryptedJWT(header, claimsSet);
      encrypt(jwt, ((LoginFromNewDeviceSettings)this.settings.get()).encryptionKey());
      return jwt.serialize();
   }

   public Result<LoginFromNewDeviceKnownUsers, LoginFromNewDeviceParsingFailure> deserialize(String rawJwt) {
      return parseEncryptedJwt(rawJwt)
         .flatMap(LoginFromNewDeviceKnownUsersSerializer::checkHeader)
         .flatMap(LoginFromNewDeviceKnownUsersSerializer::checkAlgorithm)
         .flatMap(jwt -> decrypt(((LoginFromNewDeviceSettings)this.settings.get()).encryptionKey(), jwt))
         .flatMap(LoginFromNewDeviceKnownUsersSerializer::mapToUsers);
   }

   private static Pair<String, Map<String, Object>> toJSONSerializable(Entry<UserId, LoginFromNewDeviceUserInfo> user) {
      Map<String, Object> userInfo = new HashMap<>();
      userInfo.put("created", user.getValue().createdAt().instant().getEpochSecond());
      userInfo.put("renewed", user.getValue().renewedAt().instant().getEpochSecond());
      return Pair.pair(user.getKey().getValue(), userInfo);
   }

   private static void encrypt(EncryptedJWT encryptedJWT, LoginFromNewDeviceEncryptionKey key) {
      try {
         DirectEncrypter encrypter = new DirectEncrypter(key.bytes());
         encryptedJWT.encrypt(encrypter);
      } catch (JOSEException e) {
         throw new FailureException(e);
      }
   }

   private static Result<EncryptedJWT, LoginFromNewDeviceParsingFailure> parseEncryptedJwt(String jwt) {
      try {
         EncryptedJWT parsedJWT = EncryptedJWT.parse(jwt);
         return Result.result(parsedJWT);
      } catch (ParseException e) {
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.BAD_JWT, LogUtil.logSafeQuote(jwt), e));
      }
   }

   private static Result<EncryptedJWT, LoginFromNewDeviceParsingFailure> checkHeader(EncryptedJWT jwt) {
      JWEHeader jwtHeader = jwt.getHeader();
      if (!JOSEObjectType.JWT.equals(jwtHeader.getType())) {
         String reason = String.format("was %s ; expected %s", LogUtil.quote(jwtHeader.getType()), LogUtil.quote(JOSEObjectType.JWT));
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.WRONG_TYPE_HEADER, reason));
      } else {
         return Result.result(jwt);
      }
   }

   private static Result<EncryptedJWT, LoginFromNewDeviceParsingFailure> checkAlgorithm(EncryptedJWT encryptedJWT) {
      JWEAlgorithm algorithmFromJWT = encryptedJWT.getHeader().getAlgorithm();
      if (!algorithmFromJWT.equals(DEFAULT_ALGORITHM)) {
         String reason = String.format("was %s ; expected %s", LogUtil.quote(algorithmFromJWT.getName()), LogUtil.quote(DEFAULT_ALGORITHM));
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.WRONG_ENCRYPTION_ALGORITHM, reason));
      } else {
         return Result.result(encryptedJWT);
      }
   }

   private static Result<EncryptedJWT, LoginFromNewDeviceParsingFailure> decrypt(LoginFromNewDeviceEncryptionKey encryptionKey, EncryptedJWT encryptedJWT) {
      try {
         DirectDecrypter decrypter = new DirectDecrypter(encryptionKey.bytes());
         encryptedJWT.decrypt(decrypter);
         return Result.result(encryptedJWT);
      } catch (JOSEException e) {
         String reason = "JWT might have been encrypted using a different key.";
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.DECRYPTION_ERROR, reason));
      }
   }

   private static Result<LoginFromNewDeviceKnownUsers, LoginFromNewDeviceParsingFailure> mapToUsers(EncryptedJWT encryptedJwt) {
      JWTClaimsSet claims;
      try {
         claims = encryptedJwt.getJWTClaimsSet();
      } catch (ParseException e) {
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.MAPPING_ERROR, e));
      }

      Map<String, Object> usersClaim;
      try {
         usersClaim = claims.getJSONObjectClaim("users");
      } catch (ParseException e) {
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.MAPPING_ERROR, e));
      }

      if (usersClaim == null) {
         return Result.failure(
            LoginFromNewDeviceParsingFailure.of(
               LoginFromNewDeviceParsingFailure.FailureType.MAPPING_ERROR, String.format("No claim %s in claim set", LogUtil.quote("users"))
            )
         );
      }

      LoginFromNewDeviceKnownUsers parsedUsers;
      try {
         parsedUsers = LoginFromNewDeviceKnownUsers.ofListOnOverflowDropOldest(
            usersClaim.keySet().stream().map(userId -> mapToUser(userId, usersClaim)).toList()
         );
      } catch (ViolationException e) {
         return Result.failure(LoginFromNewDeviceParsingFailure.of(LoginFromNewDeviceParsingFailure.FailureType.MAPPING_ERROR, e));
      }

      return Result.result(parsedUsers);
   }

   private static LoginFromNewDeviceUser mapToUser(String userKey, Map<String, Object> usersClaim) {
      UserId userId = UserId.userId(userKey);

      LoginFromNewDeviceCreatedAt createdAt;
      LoginFromNewDeviceRenewedAt renewedAt;
      try {
         Map<String, Object> userInfos = JSONObjectUtils.getJSONObject(usersClaim, userKey);
         createdAt = new LoginFromNewDeviceCreatedAt(DateUtils.fromSecondsSinceEpoch(JSONObjectUtils.getLong(userInfos, "created")).toInstant());
         renewedAt = new LoginFromNewDeviceRenewedAt(DateUtils.fromSecondsSinceEpoch(JSONObjectUtils.getLong(userInfos, "renewed")).toInstant());
      } catch (ParseException e) {
         throw new ViolationException(e);
      }

      return LoginFromNewDeviceUser.of(userId, createdAt, renewedAt);
   }
}