// SPDX-FileCopyrightText: 2022 Melvin Keskin <melvo@olomono.de>
// SPDX-FileCopyrightText: 2022 Linus Jahn <lnj@kaidan.im>
//
// SPDX-License-Identifier: LGPL-2.1-or-later

#ifndef QXMPPOMEMOMANAGER_P_H
#define QXMPPOMEMOMANAGER_P_H

#include "QXmppConstants_p.h"
#include "QXmppE2eeMetadata.h"
#include "QXmppOmemoDeviceBundle_p.h"
#include "QXmppOmemoManager.h"
#include "QXmppOmemoStorage.h"
#include "QXmppPubSubManager.h"

#include "Crypto.h"
#include "OmemoLibWrappers.h"

#include <QDomElement>
#include <QTimer>

#undef max

class QXmppTrustManager;
class QXmppOmemoManager;
class QXmppPubSubManager;
class QXmppPubSubNodeConfig;
class QXmppPubSubPublishOptions;
class QXmppOmemoIq;
class QXmppOmemoEnvelope;
class QXmppOmemoElement;
class QXmppOmemoDeviceListItem;
class QXmppOmemoDeviceBundleItem;

using namespace QXmpp;
using namespace std::chrono_literals;
using SecureByteArray = QXmpp::Private::Crypto::SecureByteArray;

namespace QXmpp::Omemo::Private {

// default possible trust levels a key must have to be used for encryption
// The class documentation must be adapted if the trust levels are modified.
constexpr auto ACCEPTED_TRUST_LEVELS = TrustLevel::AutomaticallyTrusted | TrustLevel::ManuallyTrusted | TrustLevel::Authenticated;

// count of unresponded stanzas sent to a device until QXmpp stops encrypting for it
constexpr int UNRESPONDED_STANZAS_UNTIL_ENCRYPTION_IS_STOPPED = 106;

// count of unresponded stanzas received from a device until a heartbeat message is sent to it
constexpr int UNRESPONDED_STANZAS_UNTIL_HEARTBEAT_MESSAGE_IS_SENT = 53;

// size of empty OMEMO message's decryption data
constexpr int EMPTY_MESSAGE_DECRYPTION_DATA_SIZE = 32;

// workaround for PubSub nodes that are not configurable to store 'max' as the value for
// 'pubsub#max_items'
constexpr uint64_t PUBSUB_NODE_MAX_ITEMS_1 = 1000;
constexpr uint64_t PUBSUB_NODE_MAX_ITEMS_2 = 100;
constexpr uint64_t PUBSUB_NODE_MAX_ITEMS_3 = 10;

constexpr uint32_t PRE_KEY_ID_MIN = 1;
constexpr uint32_t SIGNED_PRE_KEY_ID_MIN = 1;
constexpr uint32_t PRE_KEY_ID_MAX = std::numeric_limits<int32_t>::max();
constexpr uint32_t SIGNED_PRE_KEY_ID_MAX = std::numeric_limits<int32_t>::max();
constexpr uint32_t PRE_KEY_INITIAL_CREATION_COUNT = 100;

// maximum count of devices stored per JID
constexpr int DEVICES_PER_JID_MAX = 200;

// maximum count of devices for whom a stanza is encrypted
constexpr int DEVICES_PER_STANZA_MAX = 1000;

// interval to remove old signed pre keys and create new ones
constexpr auto SIGNED_PRE_KEY_RENEWAL_INTERVAL = 24h * 7 * 4;

// interval to check for old signed pre keys
constexpr auto SIGNED_PRE_KEY_RENEWAL_CHECK_INTERVAL = 24h;

// interval to remove devices locally after removal from their servers
constexpr auto DEVICE_REMOVAL_INTERVAL = 24h * 7 * 12;

// interval to check for devices removed from their servers
constexpr auto DEVICE_REMOVAL_CHECK_INTERVAL = 24h;

inline constexpr auto HKDF_INFO = "OMEMO Payload";
constexpr int HKDF_KEY_SIZE = 32;
constexpr int HKDF_SALT_SIZE = 32;
constexpr int HKDF_OUTPUT_SIZE = 80;

constexpr uint32_t PAYLOAD_MESSAGE_AUTHENTICATION_CODE_SIZE = 16;

constexpr int PAYLOAD_KEY_SIZE = 32;
constexpr uint32_t PAYLOAD_INITIALIZATION_VECTOR_SIZE = 16;
constexpr uint32_t PAYLOAD_AUTHENTICATION_KEY_SIZE = 32;

// boundaries for the count of characters in SCE's <rpad/> element
constexpr uint32_t SCE_RPAD_SIZE_MIN = 0;
constexpr uint32_t SCE_RPAD_SIZE_MAX = 200;

struct PayloadEncryptionResult {
    SecureByteArray decryptionData;
    QByteArray encryptedPayload;
};

struct DecryptionResult {
    QDomElement sceContent;
    QXmppE2eeMetadata e2eeMetadata;
};

struct IqDecryptionResult {
    QDomElement iq;
    QXmppE2eeMetadata e2eeMetadata;
};

struct OmemoEnvelopeData {
    QByteArray data;
    bool isKeyExchange;
};

}  // namespace QXmpp::Omemo::Private

using namespace QXmpp::Private;
using namespace QXmpp::Omemo::Private;

class QXmppOmemoManagerPrivate
{
public:
    using Result = std::variant<QXmpp::Success, QXmppError>;

    QXmppOmemoManager *q;

    bool initialized = false;
    bool isNewDeviceAutoSessionBuildingEnabled = false;

    QXmppOmemoStorage *omemoStorage;
    QXmppTrustManager *trustManager = nullptr;
    QXmppPubSubManager *pubSubManager = nullptr;

    QTimer signedPreKeyPairsRenewalTimer;
    QTimer deviceRemovalTimer;

    TrustLevels acceptedSessionBuildingTrustLevels = ACCEPTED_TRUST_LEVELS;

    QXmppOmemoStorage::OwnDevice ownDevice;
    QHash<uint32_t, QByteArray> preKeyPairs;
    QHash<uint32_t, QXmppOmemoStorage::SignedPreKeyPair> signedPreKeyPairs;
    QXmppOmemoDeviceBundle deviceBundle;

    int maximumDevicesPerJid = DEVICES_PER_JID_MAX;
    int maximumDevicesPerStanza = DEVICES_PER_STANZA_MAX;

    // recipient JID mapped to device ID mapped to device
    QHash<QString, QHash<uint32_t, QXmppOmemoStorage::Device>> devices;

    QList<QString> jidsOfManuallySubscribedDevices;

    OmemoContextPtr globalContext;
    StoreContextPtr storeContext;
    QRecursiveMutex mutex;
    signal_crypto_provider cryptoProvider;

    signal_protocol_identity_key_store identityKeyStore;
    signal_protocol_pre_key_store preKeyStore;
    signal_protocol_signed_pre_key_store signedPreKeyStore;
    signal_protocol_session_store sessionStore;

    QXmppOmemoManagerPrivate(QXmppOmemoManager *parent, QXmppOmemoStorage *omemoStorage);

    void initOmemoLibrary();
    // exports for unit tests
    QXMPP_EXPORT bool initGlobalContext();
    QXMPP_EXPORT bool initLocking();
    QXMPP_EXPORT bool initCryptoProvider();
    void initStores();

    signal_protocol_identity_key_store createIdentityKeyStore() const;
    signal_protocol_signed_pre_key_store createSignedPreKeyStore() const;
    signal_protocol_pre_key_store createPreKeyStore() const;
    signal_protocol_session_store createSessionStore() const;

    QXmppTask<bool> setUpDeviceId();
    std::optional<uint32_t> generateDeviceId();
    std::optional<uint32_t> generateDeviceId(const QList<QString> &existingIds);
    bool setUpIdentityKeyPair(ratchet_identity_key_pair **identityKeyPair);
    void schedulePeriodicTasks();
    void renewSignedPreKeyPairs();
    bool updateSignedPreKeyPair(ratchet_identity_key_pair *identityKeyPair);
    bool renewPreKeyPairs(uint32_t keyPairBeingRenewed);
    bool updatePreKeyPairs(uint32_t count = 1);
    void removeDevicesRemovedFromServer();

    QXmppTask<QXmppE2eeExtension::MessageEncryptResult> encryptMessageForRecipients(QXmppMessage &&message,
                                                                                    QList<QString> recipientJids,
                                                                                    TrustLevels acceptedTrustLevels);
    template<typename T>
    QXmppTask<std::optional<QXmppOmemoElement>> encryptStanza(const T &stanza, const QList<QString> &recipientJids, TrustLevels acceptedTrustLevels);
    std::optional<PayloadEncryptionResult> encryptPayload(const QByteArray &payload) const;
    template<typename T>
    QByteArray createSceEnvelope(const T &stanza);
    std::optional<OmemoEnvelopeData> createOmemoEnvelopeData(const signal_protocol_address &address, const SecureByteArray &payloadDecryptionData) const;

    QXmppTask<std::optional<QXmppMessage>> decryptMessage(QXmppMessage stanza);
    QXmppTask<std::optional<IqDecryptionResult>> decryptIq(QDomElement iqElement);
    template<typename T>
    QXmppTask<std::optional<DecryptionResult>> decryptStanza(T stanza,
                                                             QString senderJid,
                                                             uint32_t senderDeviceId,
                                                             const QXmppOmemoEnvelope &omemoEnvelope,
                                                             const QByteArray &omemoPayload,
                                                             bool isMessageStanza = true);
    QXmppTask<QByteArray> extractSceEnvelope(const QString &senderJid,
                                             uint32_t senderDeviceId,
                                             const QXmppOmemoEnvelope &omemoEnvelope,
                                             QByteArray omemoPayload,
                                             bool isMessageStanza);
    QXmppTask<std::optional<SecureByteArray>> extractPayloadDecryptionData(const QString &senderJid,
                                                                           uint32_t senderDeviceId,
                                                                           QXmppOmemoEnvelope omemoEnvelope,
                                                                           bool isMessageStanza = true);
    QByteArray decryptPayload(const SecureByteArray &payloadDecryptionData, const QByteArray &payload) const;

    QXmppTask<bool> publishOmemoData();

    QXmppTask<bool> publishDeviceBundle(bool isDeviceBundlesNodeExistent,
                                        bool arePublishOptionsSupported,
                                        bool isAutomaticCreationSupported,
                                        bool isCreationAndConfigurationSupported,
                                        bool isCreationSupported,
                                        bool isConfigurationSupported,
                                        bool isConfigNodeMaxSupported);
    QXmppTask<bool> publishDeviceBundleWithoutOptions(bool isDeviceBundlesNodeExistent,
                                                      bool isCreationAndConfigurationSupported,
                                                      bool isCreationSupported,
                                                      bool isConfigurationSupported,
                                                      bool isConfigNodeMaxSupported);
    QXmppTask<bool> configureNodeAndPublishDeviceBundle(bool isConfigNodeMaxSupported);
    QXmppTask<bool> createAndConfigureDeviceBundlesNode(bool isConfigNodeMaxSupported);
    QXmppTask<bool> createDeviceBundlesNode();
    QXmppTask<bool> configureDeviceBundlesNode(bool isConfigNodeMaxSupported);
    QXmppTask<bool> publishDeviceBundleItem();
    QXmppTask<bool> publishDeviceBundleItemWithOptions();
    QXmppOmemoDeviceBundleItem deviceBundleItem() const;
    QXmppTask<std::optional<QXmppOmemoDeviceBundle>> requestDeviceBundle(QString deviceOwnerJid, uint32_t deviceId) const;
    QXmppTask<bool> deleteDeviceBundle();

    QXmppTask<bool> publishDeviceElement(bool isDeviceListNodeExistent,
                                         bool arePublishOptionsSupported,
                                         bool isAutomaticCreationSupported,
                                         bool isCreationAndConfigurationSupported,
                                         bool isCreationSupported,
                                         bool isConfigurationSupported);
    QXmppTask<bool> publishDeviceElementWithoutOptions(bool isDeviceListNodeExistent,
                                                       bool isCreationAndConfigurationSupported,
                                                       bool isCreationSupported,
                                                       bool isConfigurationSupported);
    QXmppTask<bool> configureNodeAndPublishDeviceElement();
    QXmppTask<bool> createAndConfigureDeviceListNode();
    QXmppTask<bool> createDeviceListNode();
    QXmppTask<bool> configureDeviceListNode();
    QXmppTask<bool> publishDeviceListItem(bool addOwnDevice);
    QXmppTask<bool> publishDeviceListItemWithOptions();
    QXmppOmemoDeviceListItem deviceListItem(bool addOwnDevice = true);
    QXmppTask<bool> updateOwnDevicesLocally(bool isDeviceListNodeExistent);
    std::optional<QXmppOmemoDeviceListItem> updateContactDevices(const QString &deviceOwnerJid, const QList<QXmppOmemoDeviceListItem> &deviceListItems);
    void updateDevices(const QString &deviceOwnerJid, const QXmppOmemoDeviceListItem &deviceListItem);
    void handleIrregularDeviceListChanges(const QString &deviceOwnerJid);
    QXmppTask<bool> deleteDeviceElement();

    QXmppTask<bool> createNode(const QString &node);
    QXmppTask<bool> createNode(const QString &node, const QXmppPubSubNodeConfig &config);
    QXmppTask<bool> configureNode(const QString &node, const QXmppPubSubNodeConfig &config);
    QXmppTask<bool> retractItem(const QString &node, uint32_t itemId);
    QXmppTask<bool> deleteNode(QString node);

    template<typename T>
    QXmppTask<bool> publishItem(const QString &node, const T &item);
    template<typename T>
    QXmppTask<bool> publishItem(const QString &node, const T &item, const QXmppPubSubPublishOptions &publishOptions);

    template<typename T>
    QXmppTask<bool> runPubSubQueryWithContinuation(QXmppTask<T> future, QString errorMessage);

    QXmppTask<bool> changeDeviceLabel(QString deviceLabel);

    QXmppTask<QXmppPubSubManager::ItemResult<QXmppOmemoDeviceListItem>> requestDeviceList(QString jid);
    void subscribeToNewDeviceLists(const QString &jid, uint32_t deviceId);
    QXmppTask<Result> subscribeToDeviceList(QString jid);
    QXmppTask<QList<QXmppOmemoManager::DevicesResult>> unsubscribeFromDeviceLists(const QList<QString> &jids);
    QXmppTask<Result> unsubscribeFromDeviceList(QString jid);

    QXmppTask<bool> resetOwnDevice();
    QXmppTask<void> resetOwnDeviceLocally();
    QXmppTask<bool> resetAll();
    void resetCachedData();

    QXmppTask<bool> buildSessionForNewDevice(const QString &jid, uint32_t deviceId, QXmppOmemoStorage::Device &device);
    QXmppTask<bool> buildSessionWithDeviceBundle(QString jid, uint32_t deviceId, QXmppOmemoStorage::Device &device);
    bool buildSession(signal_protocol_address address, const QXmppOmemoDeviceBundle &deviceBundle);
    bool createSessionBundle(session_pre_key_bundle **sessionBundle,
                             const QByteArray &serializedPublicIdentityKey,
                             const QByteArray &serializedSignedPublicPreKey,
                             uint32_t signedPublicPreKeyId,
                             const QByteArray &serializedSignedPublicPreKeySignature,
                             const QByteArray &serializedPublicPreKey,
                             uint32_t publicPreKeyId);

    bool deserializeIdentityKeyPair(ratchet_identity_key_pair **identityKeyPair) const;
    bool deserializePrivateIdentityKey(ec_private_key **privateIdentityKey, const QByteArray &serializedPrivateIdentityKey) const;
    bool deserializePublicIdentityKey(ec_public_key **publicIdentityKey, const QByteArray &serializedPublicIdentityKey) const;
    bool deserializeSignedPublicPreKey(ec_public_key **signedPublicPreKey, const QByteArray &serializedSignedPublicPreKey) const;
    bool deserializePublicPreKey(ec_public_key **publicPreKey, const QByteArray &serializedPublicPreKey) const;

    QXmppTask<QXmpp::SendResult> sendEmptyMessage(const QString &recipientJid, uint32_t recipientDeviceId) const;
    QXmppTask<void> storeOwnKey() const;
    QXmppTask<TrustLevel> storeKeyDependingOnSecurityPolicy(QString keyOwnerJid, QByteArray key);
    QXmppTask<TrustLevel> storeKey(QString keyOwnerJid, QByteArray key, TrustLevel trustLevel = TrustLevel::AutomaticallyDistrusted) const;
    QString ownBareJid() const;
    QString ownFullJid() const;
    QHash<uint32_t, QXmppOmemoStorage::Device> otherOwnDevices();

    void warning(const QString &msg) const;
};

#endif  // QXMPPOMEMOMANAGER_P_H
