/**
 *
 * Copyright 2018 Paul Schaub.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.jivesoftware.smackx.ox.bouncycastle;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import org.jivesoftware.smack.util.MultiMap;
import org.jivesoftware.smackx.ox.OpenPgpProvider;
import org.jivesoftware.smackx.ox.OpenPgpV4Fingerprint;
import org.jivesoftware.smackx.ox.bouncycastle.selection_strategy.BareJidUserId;
import org.jivesoftware.smackx.ox.callback.SmackMissingOpenPgpPublicKeyCallback;
import org.jivesoftware.smackx.ox.element.CryptElement;
import org.jivesoftware.smackx.ox.element.SignElement;
import org.jivesoftware.smackx.ox.element.SigncryptElement;
import org.jivesoftware.smackx.ox.exception.MissingOpenPgpKeyPairException;
import org.jivesoftware.smackx.ox.exception.MissingOpenPgpPublicKeyException;
import org.jivesoftware.smackx.ox.exception.MissingUserIdOnKeyException;
import org.jivesoftware.smackx.ox.exception.SmackOpenPgpException;
import org.jivesoftware.smackx.ox.util.DecryptedBytesAndMetadata;
import org.jivesoftware.smackx.ox.util.KeyBytesAndFingerprint;

import de.vanitasvitae.crypto.pgpainless.PGPainless;
import de.vanitasvitae.crypto.pgpainless.decryption_verification.DecryptionStream;
import de.vanitasvitae.crypto.pgpainless.decryption_verification.MissingPublicKeyCallback;
import de.vanitasvitae.crypto.pgpainless.decryption_verification.PainlessResult;
import de.vanitasvitae.crypto.pgpainless.key.generation.type.length.RsaLength;
import de.vanitasvitae.crypto.pgpainless.util.BCUtil;
import org.bouncycastle.openpgp.PGPException;
import org.bouncycastle.openpgp.PGPPublicKey;
import org.bouncycastle.openpgp.PGPPublicKeyRing;
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection;
import org.bouncycastle.openpgp.PGPSecretKeyRing;
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection;
import org.bouncycastle.openpgp.operator.bc.BcKeyFingerprintCalculator;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.io.Streams;
import org.jxmpp.jid.BareJid;

public class PainlessOpenPgpProvider implements OpenPgpProvider {

    private final PainlessOpenPgpStore store;
    private final BareJid owner;

    public PainlessOpenPgpProvider(BareJid owner, PainlessOpenPgpStore store) {
        this.owner = owner;
        this.store = store;
    }

    @Override
    public byte[] signAndEncrypt(SigncryptElement element,
                                 OpenPgpV4Fingerprint signingKey,
                                 MultiMap<BareJid, OpenPgpV4Fingerprint> encryptionKeys)
            throws MissingOpenPgpKeyPairException, MissingOpenPgpPublicKeyException, SmackOpenPgpException,
            IOException {

        PGPPublicKeyRing[] validEncryptionKeys = getEncryptionKeys(encryptionKeys);
        PGPSecretKeyRing signingKeyRing = getSigningKey(signingKey);

        InputStream fromPlain = element.toInputStream();
        ByteArrayOutputStream encrypted = new ByteArrayOutputStream();
        OutputStream encryptor;
        try {
            encryptor = PGPainless.createEncryptor()
                    .onOutputStream(encrypted)
                    .toRecipients(validEncryptionKeys)
                    .usingSecureAlgorithms()
                    .signWith(store.getSecretKeyProtector(), signingKeyRing)
                    .noArmor();
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }

        Streams.pipeAll(fromPlain, encryptor);
        fromPlain.close();
        encryptor.close();

        return encrypted.toByteArray();
    }

    @Override
    public byte[] sign(SignElement element, OpenPgpV4Fingerprint signingKeyFingerprint)
            throws MissingOpenPgpKeyPairException, IOException, SmackOpenPgpException {
        InputStream fromPlain = element.toInputStream();
        PGPSecretKeyRing signingKeyRing;
        try {
            signingKeyRing = store.getSecretKeyRings(owner).getSecretKeyRing(signingKeyFingerprint.getKeyId());
        } catch (PGPException e) {
            throw new MissingOpenPgpKeyPairException(owner, signingKeyFingerprint, e);
        }

        ByteArrayOutputStream toSigned = new ByteArrayOutputStream();
        OutputStream signer;
        try {
            signer = PGPainless.createEncryptor().onOutputStream(toSigned)
                    .doNotEncrypt()
                    .signWith(store.getSecretKeyProtector(), signingKeyRing)
                    .noArmor();
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }

        Streams.pipeAll(fromPlain, signer);

        fromPlain.close();
        signer.close();

        return toSigned.toByteArray();
    }

    @Override
    public byte[] encrypt(CryptElement element, MultiMap<BareJid, OpenPgpV4Fingerprint> encryptionKeyFingerprints)
            throws MissingOpenPgpPublicKeyException, IOException, SmackOpenPgpException {
        PGPPublicKeyRing[] allRecipientsKeys = getEncryptionKeys(encryptionKeyFingerprints);

        InputStream fromPlain = element.toInputStream();
        ByteArrayOutputStream encrypted = new ByteArrayOutputStream();
        OutputStream encryptor;
        try {
            encryptor = PGPainless.createEncryptor()
                    .onOutputStream(encrypted)
                    .toRecipients(allRecipientsKeys)
                    .usingSecureAlgorithms()
                    .doNotSign()
                    .noArmor();
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }

        Streams.pipeAll(fromPlain, encryptor);
        fromPlain.close();
        encryptor.close();

        return encrypted.toByteArray();
    }

    @Override
    public DecryptedBytesAndMetadata decrypt(byte[] bytes, BareJid sender, final SmackMissingOpenPgpPublicKeyCallback missingPublicKeyCallback)
            throws MissingOpenPgpKeyPairException, SmackOpenPgpException, IOException {
        Set<Long> trustedKeyIds = new HashSet<>();
        Set<PGPPublicKeyRing> senderKeys = new HashSet<>();
        InputStream fromEncrypted = new ByteArrayInputStream(bytes);
        ByteArrayOutputStream toPlain = new ByteArrayOutputStream();
        DecryptionStream decryptionStream;
        try {
            decryptionStream = PGPainless.createDecryptor().onInputStream(fromEncrypted)
                    .decryptWith(store.getSecretKeyRings(owner), store.getSecretKeyProtector())
                    .verifyWith(trustedKeyIds, senderKeys)
                    .handleMissingPublicKeysWith(new MissingPublicKeyCallback() {
                        @Override
                        public void onMissingPublicKeyEncountered(Long aLong) {

                        }
                    })
                    .build();
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }
        Streams.pipeAll(decryptionStream, toPlain);
        fromEncrypted.close();
        decryptionStream.close();

        PainlessResult result = decryptionStream.getResult();

        return new DecryptedBytesAndMetadata(toPlain.toByteArray(),
                result.getVerifiedSignatureKeyIds(),
                result.getDecryptionKeyId());
    }

    @Override
    public byte[] symmetricallyEncryptWithPassword(byte[] bytes, String password)
            throws SmackOpenPgpException, IOException {
        try {
            return PGPainless.encryptWithPassword(bytes, password.toCharArray());
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }
    }

    @Override
    public byte[] symmetricallyDecryptWithPassword(byte[] bytes, String password)
            throws SmackOpenPgpException, IOException {
        try {
            return PGPainless.decryptWithPassword(bytes, password.toCharArray());
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }
    }

    @Override
    public PainlessOpenPgpStore getStore() {
        return store;
    }

    private PGPPublicKeyRing[] getEncryptionKeys(MultiMap<BareJid, OpenPgpV4Fingerprint> encryptionKeys)
            throws IOException, SmackOpenPgpException {
        Set<PGPPublicKeyRing> allRecipientsKeys = new HashSet<>();

        for (BareJid recipient : encryptionKeys.keySet()) {
            PGPPublicKeyRingCollection recipientsKeyRings;
            try {
                recipientsKeyRings = store.getPublicKeyRings(recipient);
                for (OpenPgpV4Fingerprint fingerprint : encryptionKeys.getAll(recipient)) {
                    PGPPublicKeyRing ring = recipientsKeyRings.getPublicKeyRing(fingerprint.getKeyId());
                    if (ring != null) allRecipientsKeys.add(ring);
                }
            } catch (PGPException e) {
                throw new SmackOpenPgpException(e);
            }
        }

        PGPPublicKeyRing[] allEncryptionKeys = new PGPPublicKeyRing[allRecipientsKeys.size()];
        Iterator<PGPPublicKeyRing> iterator = allRecipientsKeys.iterator();
        for (int i = 0; i < allEncryptionKeys.length; i++) {
            allEncryptionKeys[i] = iterator.next();
        }

        return allEncryptionKeys;
    }

    private PGPSecretKeyRing getSigningKey(OpenPgpV4Fingerprint signingKey)
            throws IOException, MissingOpenPgpKeyPairException {
        PGPSecretKeyRing signingKeyRing;
        try {
            signingKeyRing = store.getSecretKeyRings(owner).getSecretKeyRing(signingKey.getKeyId());
        } catch (PGPException e) {
            throw new MissingOpenPgpKeyPairException(owner, signingKey, e);
        }
        return signingKeyRing;
    }

    @Override
    public KeyBytesAndFingerprint generateOpenPgpKeyPair(BareJid owner)
            throws SmackOpenPgpException, InvalidAlgorithmParameterException, NoSuchAlgorithmException,
            NoSuchProviderException, IOException {
        PGPSecretKeyRing secretKey;
        try {
            secretKey = PGPainless.generateKeyRing().simpleRsaKeyRing("xmpp:" + owner.toString(), RsaLength._4096);
        } catch (PGPException e) {
            throw new SmackOpenPgpException("Could not generate OpenPGP Key Pair.", e);
        }

        return new KeyBytesAndFingerprint(secretKey.getEncoded(), getFingerprint(secretKey.getPublicKey()));
    }

    @Override
    public OpenPgpV4Fingerprint importPublicKey(BareJid owner, byte[] bytes)
            throws MissingUserIdOnKeyException, IOException, SmackOpenPgpException {
        PGPPublicKeyRing publicKeys = new PGPPublicKeyRing(bytes, new BcKeyFingerprintCalculator());
        return importPublicKey(owner, publicKeys);
    }

    public OpenPgpV4Fingerprint importPublicKey(BareJid owner, PGPPublicKeyRing ring)
            throws SmackOpenPgpException, IOException, MissingUserIdOnKeyException {
        if (!new BareJidUserId.PubRingSelectionStrategy().accept(owner, ring)) {
            throw new MissingUserIdOnKeyException(owner, ring.getPublicKey().getKeyID());
        }
        try {
            PGPPublicKeyRingCollection publicKeyRings = getStore().getPublicKeyRings(owner);
            if (publicKeyRings == null) {
                publicKeyRings = new PGPPublicKeyRingCollection(Collections.singleton(ring));
            } else {
                publicKeyRings = PGPPublicKeyRingCollection.addPublicKeyRing(publicKeyRings, ring);
            }
            getStore().storePublicKeyRing(owner, publicKeyRings);
        } catch (PGPException e) {
            throw new SmackOpenPgpException(e);
        }
        return getFingerprint(ring.getPublicKey());
    }

    @Override
    public OpenPgpV4Fingerprint importSecretKey(BareJid owner, byte[] bytes)
            throws MissingUserIdOnKeyException, SmackOpenPgpException, IOException {
        PGPSecretKeyRing secretKeys;
        try {
            secretKeys = new PGPSecretKeyRing(bytes, new BcKeyFingerprintCalculator());
        } catch (PGPException | IOException e) {
            throw new SmackOpenPgpException("Could not deserialize PGP secret key of " + owner.toString(), e);
        }

        if (!new BareJidUserId.SecRingSelectionStrategy().accept(owner, secretKeys)) {
            throw new MissingUserIdOnKeyException(owner, secretKeys.getPublicKey().getKeyID());
        }

        PGPSecretKeyRingCollection secretKeyRings;
        try {
            secretKeyRings = getStore().getSecretKeyRings(owner);
        } catch (PGPException | IOException e) {
            throw new SmackOpenPgpException("Could not load secret key ring of " + owner.toString(), e);
        }
        if (secretKeyRings == null) {
            try {
                secretKeyRings = new PGPSecretKeyRingCollection(Collections.singleton(secretKeys));
            } catch (IOException | PGPException e) {
                throw new SmackOpenPgpException("Could not create SecretKeyRingCollection from SecretKeyRing.", e);
            }
        } else {
            secretKeyRings = PGPSecretKeyRingCollection.addSecretKeyRing(secretKeyRings, secretKeys);
        }
        getStore().storeSecretKeyRing(owner, secretKeyRings);

        PGPPublicKeyRing publicKeys = BCUtil.publicKeyRingFromSecretKeyRing(secretKeys);
        importPublicKey(owner, publicKeys);

        return getFingerprint(publicKeys.getPublicKey());
    }

    public static OpenPgpV4Fingerprint getFingerprint(PGPPublicKey publicKey) {
        byte[] hex = Hex.encode(publicKey.getFingerprint());
        return new OpenPgpV4Fingerprint(hex);
    }
}