package org.jivesoftware.smackx.ikey;

import org.bouncycastle.openpgp.PGPException;
import org.bouncycastle.openpgp.PGPPublicKeyRing;
import org.bouncycastle.openpgp.PGPSecretKeyRing;
import org.bouncycastle.util.io.Streams;
import org.jivesoftware.smack.Manager;
import org.jivesoftware.smack.SmackException;
import org.jivesoftware.smack.XMPPConnection;
import org.jivesoftware.smack.XMPPException;
import org.jivesoftware.smack.packet.Message;
import org.jivesoftware.smack.provider.ProviderManager;
import org.jivesoftware.smackx.ikey.element.IkeyElement;
import org.jivesoftware.smackx.ikey.element.ProofElement;
import org.jivesoftware.smackx.ikey.element.SignedElement;
import org.jivesoftware.smackx.ikey.element.SubordinateElement;
import org.jivesoftware.smackx.ikey.element.SubordinateListElement;
import org.jivesoftware.smackx.ikey.element.SuperordinateElement;
import org.jivesoftware.smackx.ikey.mechanism.IkeySignatureCreationMechanism;
import org.jivesoftware.smackx.ikey.mechanism.IkeySignatureVerificationMechanism;
import org.jivesoftware.smackx.ikey.provider.IkeyElementProvider;
import org.jivesoftware.smackx.ikey.provider.SubordinateListElementProvider;
import org.jivesoftware.smackx.ikey.record.IkeyRecord;
import org.jivesoftware.smackx.ikey.record.IkeyStore;
import org.jivesoftware.smackx.ikey.record.IkeySubordinateRecord;
import org.jivesoftware.smackx.ikey.record.OxSubordinateRecord;
import org.jivesoftware.smackx.ikey.util.IkeyConstants;
import org.jivesoftware.smackx.ikey.util.UnsupportedSignatureAlgorithmException;
import org.jivesoftware.smackx.ikey_ox.OxIkeySignatureCreationMechanism;
import org.jivesoftware.smackx.ikey_ox.OxIkeySignatureVerificationMechanism;
import org.jivesoftware.smackx.ox.OpenPgpManager;
import org.jivesoftware.smackx.ox.OpenPgpSecretKeyBackupPassphrase;
import org.jivesoftware.smackx.ox.element.OpenPgpElement;
import org.jivesoftware.smackx.ox.element.SecretkeyElement;
import org.jivesoftware.smackx.ox.exception.InvalidBackupCodeException;
import org.jivesoftware.smackx.ox.store.definition.OpenPgpTrustStore;
import org.jivesoftware.smackx.ox.util.OpenPgpPubSubUtil;
import org.jivesoftware.smackx.ox.util.SecretKeyBackupHelper;
import org.jivesoftware.smackx.pep.PepEventListener;
import org.jivesoftware.smackx.pep.PepManager;
import org.jivesoftware.smackx.pubsub.LeafNode;
import org.jivesoftware.smackx.pubsub.PayloadItem;
import org.jivesoftware.smackx.pubsub.PubSubException;
import org.jivesoftware.smackx.pubsub.PubSubManager;
import org.jivesoftware.smackx.pubsub.PubSubUri;
import org.jxmpp.jid.EntityBareJid;
import org.mercury_im.messenger.core.crypto.OpenPgpSecretKeyBackupPassphraseGenerator;
import org.mercury_im.messenger.core.crypto.SecureRandomSecretKeyBackupPassphraseGenerator;
import org.pgpainless.PGPainless;
import org.pgpainless.decryption_verification.DecryptionStream;
import org.pgpainless.key.OpenPgpV4Fingerprint;
import org.pgpainless.key.collection.PGPKeyRing;
import org.pgpainless.key.protection.SecretKeyRingProtector;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.WeakHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

import static org.jivesoftware.smackx.ikey.util.IkeyConstants.SUPERORDINATE_NODE;

public final class IkeyManager extends Manager {

    private static final Logger LOGGER = Logger.getLogger(IkeyManager.class.getName());
    private static final Map<XMPPConnection, IkeyManager> INSTANCES = new WeakHashMap<>();

    private final OpenPgpSecretKeyBackupPassphraseGenerator backupPassphraseGenerator;

    static {
        // TODO: Replace with .providers file once merged into Smack
        ProviderManager.addExtensionProvider(IkeyElement.ELEMENT, IkeyElement.NAMESPACE, new IkeyElementProvider());
        ProviderManager.addExtensionProvider(SubordinateListElement.ELEMENT, SubordinateListElement.NAMESPACE, new SubordinateListElementProvider());
    }

    private IkeyStore store;

    private IkeyManager(XMPPConnection connection) {
        super(connection);
        this.backupPassphraseGenerator = new SecureRandomSecretKeyBackupPassphraseGenerator();
    }

    public static synchronized IkeyManager getInstanceFor(XMPPConnection connection) {
        IkeyManager manager = INSTANCES.get(connection);
        if (manager == null) {
            manager = new IkeyManager(connection);
            INSTANCES.put(connection, manager);
        }
        return manager;
    }

    public SecretkeyElement fetchSecretIdentityKey()
            throws InterruptedException, PubSubException.NotALeafNodeException,
            XMPPException.XMPPErrorException, SmackException.NotConnectedException,
            SmackException.NoResponseException {
        return OpenPgpPubSubUtil.fetchSecretKey(PepManager.getInstanceFor(connection()), SUPERORDINATE_NODE);
    }

    public OpenPgpSecretKeyBackupPassphrase depositIdentityKeyBackup()
            throws PGPException, InterruptedException, SmackException.NoResponseException,
            SmackException.NotConnectedException, SmackException.FeatureNotSupportedException,
            XMPPException.XMPPErrorException, PubSubException.NotALeafNodeException, IOException {
        PGPSecretKeyRing secretKeys = store.loadSecretKey();
        OpenPgpSecretKeyBackupPassphrase passphrase = store.loadBackupPassphrase();

        return depositIdentityKeyBackup(secretKeys, passphrase);
    }

    public OpenPgpSecretKeyBackupPassphrase depositIdentityKeyBackup(PGPSecretKeyRing secretKey, OpenPgpSecretKeyBackupPassphrase passphrase)
            throws InterruptedException, SmackException.NoResponseException,
            SmackException.NotConnectedException, SmackException.FeatureNotSupportedException,
            XMPPException.XMPPErrorException, PubSubException.NotALeafNodeException, IOException, PGPException {
        SecretkeyElement secretkeyElement = SecretKeyBackupHelper.createSecretkeyElement(secretKey.getEncoded(), passphrase);
        OpenPgpPubSubUtil.depositSecretKey(connection(), secretkeyElement, SUPERORDINATE_NODE);

        return passphrase;
    }

    public IkeyElement createOxIkeyElement(PGPSecretKeyRing secretKeys,
                                           SecretKeyRingProtector keyRingProtector,
                                           SubordinateElement... subordinateElements) throws IOException {
        IkeySignatureCreationMechanism mechanism = new OxIkeySignatureCreationMechanism(secretKeys, keyRingProtector);
        SuperordinateElement superordinateElement = new SuperordinateElement(secretKeys.getPublicKey().getEncoded());
        SubordinateListElement subordinateListElement = new SubordinateListElement(connection().getUser().asEntityBareJid(),
                new Date(), Arrays.asList(subordinateElements));
        return createIkeyElement(mechanism, superordinateElement, subordinateListElement);
    }

    public boolean deleteSecretIdentityKeyNode()
            throws XMPPException.XMPPErrorException, SmackException.NotConnectedException,
            InterruptedException, SmackException.NoResponseException {
        return OpenPgpPubSubUtil.deleteSecretKeyNode(PepManager.getInstanceFor(connection()), SUPERORDINATE_NODE);
    }

    public void startListeners() {
        PepManager.getInstanceFor(connection())
                .addPepEventListener(IkeyConstants.SUBORDINATES_NODE, IkeyElement.class, ikeyPepEventListener);
    }

    public void stopListeners() {
        PepManager.getInstanceFor(connection())
                .removePepEventListener(ikeyPepEventListener);
    }

    public void setStore(IkeyStore store) {
        this.store = store;
    }

    public IkeyElement createIkeyElement(IkeySignatureCreationMechanism mechanism,
                                         SuperordinateElement superordinateElement,
                                         SubordinateListElement subordinateListElement)
            throws IOException {
        SignedElement signedElement = new SignedElement(subordinateListElement);
        ProofElement proofElement = new ProofElement(mechanism.createSignature(signedElement.getUtf8Bytes()));
        return new IkeyElement(mechanism.getType(), superordinateElement, signedElement, proofElement);
    }

    public void storeAndPublishElement(IkeyElement ikeyElement)
            throws InterruptedException, PubSubException.NotALeafNodeException, XMPPException.XMPPErrorException,
            SmackException.NotConnectedException, SmackException.NoResponseException, IOException, PGPException {
        IkeyRecord record = elementToRecord(ikeyElement);
        record.setTrust(OpenPgpTrustStore.Trust.trusted);
        store.storeIkeyRecord(connection().getUser().asEntityBareJid(), record);
        publishIkeyElement(ikeyElement);
    }

    public void publishIkeyElement(IkeyElement element)
            throws InterruptedException, PubSubException.NotALeafNodeException,
            XMPPException.XMPPErrorException, SmackException.NotConnectedException,
            SmackException.NoResponseException {
        PepManager.getInstanceFor(connection())
                .publish(IkeyConstants.SUBORDINATES_NODE, new PayloadItem<>(element));
    }

    public IkeyElement fetchIkeyElementOf(EntityBareJid jid)
            throws InterruptedException, PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, XMPPException.XMPPErrorException, PubSubException.NotAPubSubNodeException {
        PubSubManager pubSubManager = PubSubManager.getInstanceFor(connection(), jid);
        return fetchIkeyElementFrom(pubSubManager);
    }

    public IkeyElement fetchOwnIkeyElement()
            throws InterruptedException, PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, XMPPException.XMPPErrorException, PubSubException.NotAPubSubNodeException {
        PubSubManager pubSubManager = PubSubManager.getInstanceFor(connection());
        return fetchIkeyElementFrom(pubSubManager);
    }

    private static IkeyElement fetchIkeyElementFrom(PubSubManager pubSubManager)
            throws PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, InterruptedException, XMPPException.XMPPErrorException,
            PubSubException.NotAPubSubNodeException {
        LeafNode node = pubSubManager.getLeafNode(IkeyConstants.SUBORDINATES_NODE);
        List<PayloadItem<IkeyElement>> items = node.getItems(1);
        if (items.isEmpty()) {
            return null;
        } else {
            return items.get(0).getPayload();
        }
    }

    private void processIkeyElement(EntityBareJid from, IkeyElement element)
            throws IOException, UnsupportedSignatureAlgorithmException, PGPException {
        IkeyRecord newRecord = elementToRecord(element);
        if (isFromTheFuture(newRecord)) {
            LOGGER.log(Level.WARNING, "Received ikey element appears to be from the future: " +
                    newRecord.getTimestamp());
            return;
        }

        IkeyRecord existingRecord = store.loadIkeyRecord(newRecord.getJid());

        if (existsSameOrNewerRecord(newRecord, existingRecord)) {
            LOGGER.log(Level.WARNING, "There exists this exact, or a newer ikey record in the database for " + from);
            return;
        }

        if (!verifyIkeyElement(from, element)) {
            LOGGER.log(Level.WARNING, "Invalid signature on ikey element of " + from);
            return;
        }

        if (isContenderElement(newRecord, store.loadIkeyRecord(from))) {
            LOGGER.log(Level.INFO, "Storing contender element for " + from);
            store.storeContenderIkeyRecord(from, newRecord);
        } else {
            if (existingRecord != null) {
                newRecord.setTrust(existingRecord.getTrust());
            }
            PGPSecretKeyRing secretKeys = store.loadSecretKey();
            if (secretKeys != null && new OpenPgpV4Fingerprint(secretKeys).equals(new OpenPgpV4Fingerprint(newRecord.getSuperordinate()))) {
                newRecord.setTrust(OpenPgpTrustStore.Trust.trusted);
            }
            LOGGER.log(Level.INFO, "Storing ikey record " + newRecord);
            store.storeIkeyRecord(from, newRecord);
        }
    }

    private IkeyRecord elementToRecord(IkeyElement element) throws IOException, PGPException {
        SubordinateListElement childElement = element.getSignedElement().getChildElement();

        EntityBareJid jid = childElement.getJid().asEntityBareJidOrThrow();
        Date timestamp = childElement.getTimestamp();

        byte[] pubKeyBytes = element.getSuperordinate().getPubKeyBytes();
        PGPPublicKeyRing publicKeys = PGPainless.readKeyRing().publicKeyRing(pubKeyBytes);
        OpenPgpV4Fingerprint fingerprint = new OpenPgpV4Fingerprint(publicKeys);

        ByteArrayInputStream signed = new ByteArrayInputStream(element.getSignedElement().getUtf8Bytes());
        DecryptionStream verifier = PGPainless.createDecryptor().onInputStream(signed)
                .doNotDecrypt().verifyDetachedSignature(element.getProof().getSignatureBytes())
                .verifyWith(publicKeys)
                .ignoreMissingPublicKeys()
                .build();

        ByteArrayOutputStream dummyOut = new ByteArrayOutputStream();
        Streams.pipeAll(verifier, dummyOut);
        verifier.close();
        dummyOut.close();

        if (!verifier.getResult().getVerifiedSignatures().containsKey(fingerprint)) {
            throw new PGPException("Signature could not be verified.");
        }

        List<IkeySubordinateRecord> subordinateRecords = new ArrayList<>();
        for (SubordinateElement subordinate : element.getSignedElement().getChildElement().getSubordinates()) {
            if (subordinate.getType().equals(OpenPgpElement.NAMESPACE)) {
                OxSubordinateRecord sr = new OxSubordinateRecord();
                sr.setUri(subordinate.getUri());
                sr.setOxFingerprint(new OpenPgpV4Fingerprint(subordinate.getFingerprint()));
                subordinateRecords.add(sr);
            }
        }
        IkeyRecord record = new IkeyRecord();
        record.setJid(jid);
        record.setTimestamp(timestamp);
        record.setSuperordinate(publicKeys);
        record.getSubordinates().addAll(subordinateRecords);
        return record;
    }

    private boolean isContenderElement(IkeyRecord potentialContenderElement, IkeyRecord record) {
        if (record == null) {
            return false;
        }
        PGPPublicKeyRing potentialContenderKeys = potentialContenderElement.getSuperordinate();
        PGPPublicKeyRing previousKeys = record.getSuperordinate();

        return potentialContenderKeys.getPublicKey().getKeyID() != previousKeys.getPublicKey().getKeyID();
    }

    public IkeyRecord getIkeyElementOf(EntityBareJid from)
            throws IOException, InterruptedException, PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, XMPPException.XMPPErrorException, PubSubException.NotAPubSubNodeException, PGPException {
        IkeyRecord stored = store.loadIkeyRecord(from);
        if (stored == null) {
            IkeyElement element = fetchIkeyElementOf(from);
            stored = elementToRecord(element);
            store.storeIkeyRecord(from, stored);
        }
        return stored;
    }

    private boolean verifyIkeyElement(EntityBareJid from, IkeyElement element)
            throws IOException, UnsupportedSignatureAlgorithmException {
        IkeySignatureVerificationMechanism verificationMechanism = getSignatureVerificationMechanismFor(element);
        IkeySignatureVerifier verifier = new IkeySignatureVerifier(verificationMechanism);
        return verifier.verify(element, from);
    }

    private static IkeySignatureVerificationMechanism getSignatureVerificationMechanismFor(IkeyElement ikeyElement)
            throws IOException, UnsupportedSignatureAlgorithmException {
        switch (ikeyElement.getType()) {
            case OX:
                PGPPublicKeyRing ikey = PGPainless.readKeyRing().publicKeyRing(ikeyElement.getSuperordinate().getPubKeyBytes());
                return new OxIkeySignatureVerificationMechanism(ikey);

            default:
                throw new UnsupportedSignatureAlgorithmException(ikeyElement.getType().name());
        }
    }

    private static boolean isFromTheFuture(IkeyRecord record) {
        Date timestamp = record.getTimestamp();
        Date now = new Date();
        return timestamp.after(now);
    }

    private boolean existsSameOrNewerRecord(IkeyRecord record, IkeyRecord existingRecord) throws IOException {
        if (existingRecord == null) {
            return false;
        }
        Date latestTimestamp = existingRecord.getTimestamp();
        Date eventTimestamp = record.getTimestamp();
        return latestTimestamp.equals(eventTimestamp) // same
                || latestTimestamp.after(eventTimestamp); // newer
    }

    @SuppressWarnings("UnnecessaryAnonymousClass")
    private final PepEventListener<IkeyElement> ikeyPepEventListener = new PepEventListener<IkeyElement>() {
        @Override
        public void onPepEvent(EntityBareJid from, IkeyElement event, String id, Message carrierMessage) {
            try {
                processIkeyElement(from, event);
            } catch (IOException | UnsupportedSignatureAlgorithmException e) {
                LOGGER.log(Level.WARNING, "Error:", e);
            } catch (PGPException e) {
                e.printStackTrace();
            }
        }
    };

    public boolean hasStore() {
        return store != null;
    }

    public boolean hasLocalKey() {
        return store.loadSecretKey() != null;
    }

    public void generateIdentityKey() throws PGPException, NoSuchAlgorithmException, NoSuchProviderException, InvalidAlgorithmParameterException {
        PGPKeyRing key = OpenPgpManager.getInstanceFor(connection())
                .generateKeyRing(connection().getUser().asBareJid());
        store.storeSecretKey(key.getSecretKeys());
        store.storeBackupPassphrase(generateBackupPassphrase());
    }

    private OpenPgpSecretKeyBackupPassphrase generateBackupPassphrase() {
        return backupPassphraseGenerator.generateBackupPassphrase();
    }

    public void restoreSecretKeyBackup(SecretkeyElement secretkeyElement, OpenPgpSecretKeyBackupPassphrase passphrase)
            throws PGPException, IOException, InvalidBackupCodeException {
        PGPSecretKeyRing secretKeys = SecretKeyBackupHelper.restoreSecretKeyBackup(secretkeyElement, passphrase);
        store.storeSecretKey(secretKeys);
        store.storeBackupPassphrase(passphrase);
    }
}