package org.sagebionetworks.bridge.services;

import static com.google.common.base.Preconditions.checkNotNull;

import org.sagebionetworks.bridge.BridgeConstants;
import org.sagebionetworks.bridge.crypto.AesGcmEncryptor;
import org.sagebionetworks.bridge.crypto.BridgeEncryptor;
import org.sagebionetworks.bridge.models.HealthId;
import org.sagebionetworks.bridge.models.studies.Study;

import com.stormpath.sdk.account.Account;
import com.stormpath.sdk.directory.CustomData;

public class AccountEncryptionServiceImpl implements AccountEncryptionService {

    private static final int CURRENT_VERSION = 2;
    private BridgeEncryptor healthCodeEncryptorOld;
    private AesGcmEncryptor healthCodeEncryptor;
    private HealthCodeService healthCodeService;

    public void setHealthCodeEncryptorOld(BridgeEncryptor encryptor) {
        this.healthCodeEncryptorOld = encryptor;
    }

    public void setHealthCodeEncryptor(AesGcmEncryptor encryptor) {
        this.healthCodeEncryptor = encryptor;
    }

    public void setHealthCodeService(HealthCodeService healthCodeService) {
        this.healthCodeService = healthCodeService;
    }

    @Override
    public HealthId createAndSaveHealthCode(Study study, Account account) {
        checkNotNull(study);
        checkNotNull(account);
        final CustomData customData = account.getCustomData();
        final HealthId healthId = healthCodeService.create();
        final String encryptedHealthId = healthCodeEncryptor.encrypt(healthId.getId());
        customData.put(getHealthIdKey(study), encryptedHealthId);
        customData.put(getVersionKey(study), CURRENT_VERSION);
        customData.save();
        return healthId;
    }

    @Override
    public HealthId getHealthCode(Study study, Account account) {
        checkNotNull(study);
        checkNotNull(account);
        CustomData customData = account.getCustomData();
        Object healthIdObj = customData.get(getHealthIdKey(study));
        Object versionObj = customData.get(getVersionKey(study));
        if (versionObj == null) {
            // TODO: Remove me after backfill
            versionObj = customData.get(BridgeConstants.CUSTOM_DATA_VERSION);
        }
        return (getHealthId(healthIdObj, versionObj));
    }

    private String getHealthIdKey(Study study) {
        return study.getIdentifier() + BridgeConstants.CUSTOM_DATA_HEALTH_CODE_SUFFIX;
    }

    private String getVersionKey(Study study) {
        return study.getIdentifier() + BridgeConstants.CUSTOM_DATA_VERSION;
    }

    private HealthId getHealthId(final Object healthIdObj, final Object versionObj) {
        if (healthIdObj == null) {
            return null;
        }
        final int version = versionObj == null ? 1 : (Integer)versionObj;
        String hid = null;
        try {
            hid = version == CURRENT_VERSION ?
                    healthCodeEncryptor.decrypt((String) healthIdObj) :
                    healthCodeEncryptorOld.decrypt((String) healthIdObj);
        } catch(Exception e) {
            // TODO: Remove me
            hid = healthCodeEncryptorOld.decrypt((String) healthIdObj);
        }
        if (hid == null) {
            return null;
        }
        final String healthId = hid;
        final String healthCode = healthCodeService.getHealthCode(healthId);
        return new HealthId() {
            @Override
            public String getId() {
                return healthId;
            }
            @Override
            public String getCode() {
                return healthCode;
            }
        };
    }
}

@@ -378,16 +378,73 @@
 s.Study;
+%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;
 %0A%0Aimport
@@ -608,16 +608,112 @@
 vice %7B%0A%0A
+    private final Logger logger = LoggerFactory.getLogger(AccountEncryptionServiceImpl.class);%0A%0A
     priv
@@ -2822,32 +2822,136 @@
 null;%0A        %7D%0A
+%0A        logger.info(%22healthIdObj: %22 + healthIdObj);%0A        logger.info(%22versionObj: %22 + versionObj);%0A%0A
         final in
@@ -3031,24 +3031,70 @@
 hid = null;%0A
+        if (version == CURRENT_VERSION) %7B%0A    
         try 
@@ -3111,62 +3111,170 @@
     
+    
 hid = 
-version == CURRENT_VERSION ?%0A                   
+healthCodeEncryptor.decrypt((String) healthIdObj);%0A            %7D catch(Exception e) %7B%0A                logger.info(e.getMessage());%0A                hid =
  hea
@@ -3281,32 +3281,35 @@
 lthCodeEncryptor
+Old
 .decrypt((String
@@ -3326,30 +3326,80 @@
 Obj)
- :%0A                   
+;%0A            %7D%0A        %7D else %7B%0A            try %7B%0A                hid =
  hea
@@ -3449,32 +3449,36 @@
 IdObj);%0A        
+    
 %7D catch(Exceptio
@@ -3500,27 +3500,45 @@
     
-// TODO: Remove me%0A
+    logger.info(e.getMessage());%0A    
     
@@ -3597,32 +3597,46 @@
 ) healthIdObj);%0A
+            %7D%0A
         %7D%0A      

Patch develop