// 
// Decompiled by Procyon v0.6.0
// 

package com.google.crypto.tink.prf;

import com.google.crypto.tink.prf.internal.LegacyFullPrf;
import com.google.crypto.tink.internal.PrimitiveRegistry;
import com.google.crypto.tink.internal.MutablePrimitiveRegistry;
import com.google.crypto.tink.internal.MonitoringClient;
import java.util.Map;
import java.security.GeneralSecurityException;
import com.google.crypto.tink.KeyStatus;
import java.util.HashMap;
import com.google.crypto.tink.internal.MonitoringUtil;
import com.google.crypto.tink.internal.MutableMonitoringRegistry;
import com.google.crypto.tink.internal.MonitoringAnnotations;
import com.google.crypto.tink.internal.KeysetHandleInterface;
import com.google.crypto.tink.internal.LegacyProtoKey;
import com.google.crypto.tink.internal.PrimitiveConstructor;
import com.google.errorprone.annotations.Immutable;
import com.google.crypto.tink.internal.PrimitiveWrapper;

@Immutable
public class PrfSetWrapper implements PrimitiveWrapper<Prf, PrfSet>
{
    private static final PrfSetWrapper WRAPPER;
    private static final PrimitiveConstructor<LegacyProtoKey, Prf> LEGACY_FULL_PRF_PRIMITIVE_CONSTRUCTOR;
    
    @Override
    public PrfSet wrap(final KeysetHandleInterface keysetHandle, final MonitoringAnnotations annotations, final PrimitiveFactory<Prf> factory) throws GeneralSecurityException {
        MonitoringClient.Logger logger;
        if (!annotations.isEmpty()) {
            final MonitoringClient client = MutableMonitoringRegistry.globalInstance().getMonitoringClient();
            logger = client.createLogger(keysetHandle, annotations, "prf", "compute");
        }
        else {
            logger = MonitoringUtil.DO_NOTHING_LOGGER;
        }
        final Map<Integer, Prf> mutablePrfMap = new HashMap<Integer, Prf>();
        for (int i = 0; i < keysetHandle.size(); ++i) {
            final KeysetHandleInterface.Entry entry = keysetHandle.getAt(i);
            if (entry.getStatus().equals(KeyStatus.ENABLED)) {
                if (entry.getKey() instanceof LegacyProtoKey) {
                    final LegacyProtoKey legacyProtoKey = (LegacyProtoKey)entry.getKey();
                    if (legacyProtoKey.getOutputPrefix().size() != 0) {
                        throw new GeneralSecurityException("Cannot build PrfSet with keys with non-empty output prefix");
                    }
                }
                final Prf prf = factory.create(entry);
                mutablePrfMap.put(entry.getId(), new WrappedPrfSet.PrfWithMonitoring(prf, entry.getId(), logger));
            }
        }
        return new WrappedPrfSet((Map)mutablePrfMap, keysetHandle.getPrimary().getId());
    }
    
    @Override
    public Class<PrfSet> getPrimitiveClass() {
        return PrfSet.class;
    }
    
    @Override
    public Class<Prf> getInputPrimitiveClass() {
        return Prf.class;
    }
    
    public static void register() throws GeneralSecurityException {
        MutablePrimitiveRegistry.globalInstance().registerPrimitiveWrapper((PrimitiveWrapper<Object, Object>)PrfSetWrapper.WRAPPER);
        MutablePrimitiveRegistry.globalInstance().registerPrimitiveConstructor(PrfSetWrapper.LEGACY_FULL_PRF_PRIMITIVE_CONSTRUCTOR);
    }
    
    public static void registerToInternalPrimitiveRegistry(final PrimitiveRegistry.Builder primitiveRegistryBuilder) throws GeneralSecurityException {
        primitiveRegistryBuilder.registerPrimitiveWrapper((PrimitiveWrapper<Object, Object>)PrfSetWrapper.WRAPPER);
    }
    
    static {
        WRAPPER = new PrfSetWrapper();
        LEGACY_FULL_PRF_PRIMITIVE_CONSTRUCTOR = PrimitiveConstructor.create(LegacyFullPrf::create, LegacyProtoKey.class, Prf.class);
    }
    
    private static class WrappedPrfSet extends PrfSet
    {
        private final Map<Integer, Prf> keyIdToPrfMap;
        private final int primaryKeyId;
        
        private WrappedPrfSet(final Map<Integer, Prf> keyIdToPrfMap, final int primaryKeyId) {
            this.keyIdToPrfMap = keyIdToPrfMap;
            this.primaryKeyId = primaryKeyId;
        }
        
        @Override
        public int getPrimaryId() {
            return this.primaryKeyId;
        }
        
        @Override
        public Map<Integer, Prf> getPrfs() throws GeneralSecurityException {
            return this.keyIdToPrfMap;
        }
        
        @Immutable
        private static class PrfWithMonitoring implements Prf
        {
            private final Prf prf;
            private final int keyId;
            private final MonitoringClient.Logger logger;
            
            @Override
            public byte[] compute(final byte[] input, final int outputLength) throws GeneralSecurityException {
                try {
                    final byte[] output = this.prf.compute(input, outputLength);
                    this.logger.log(this.keyId, input.length);
                    return output;
                }
                catch (final GeneralSecurityException e) {
                    this.logger.logFailure();
                    throw e;
                }
            }
            
            public PrfWithMonitoring(final Prf prf, final int keyId, final MonitoringClient.Logger logger) {
                this.prf = prf;
                this.keyId = keyId;
                this.logger = logger;
            }
        }
    }
}