// DragonX Wallet - HushChat crypto primitives (implementation).

#include "chat_crypto.h"

#include <sodium.h>

#include <cstring>
#include <vector>

namespace dragonx::chat {

namespace {

// Local constants tied to the libsodium primitive. (The dev-only chat_fixture_tooling.h
// declares equivalents, but that header is not linked into the app.)
constexpr std::size_t kStreamHeaderBytes = 24;  // crypto_secretstream_xchacha20poly1305_HEADERBYTES
constexpr std::size_t kStreamABytes = 17;       // crypto_secretstream_xchacha20poly1305_ABYTES

static_assert(kChatKeyBytes == crypto_kx_PUBLICKEYBYTES, "kx public key size mismatch");
static_assert(kChatKeyBytes == crypto_kx_SECRETKEYBYTES, "kx secret key size mismatch");

// Decode exactly outLen bytes from a lowercase/uppercase hex string; reject any other length.
bool hexToFixed(const std::string& hex, unsigned char* out, std::size_t outLen) {
    if (hex.size() != outLen * 2) return false;
    std::size_t binLen = 0;
    if (sodium_hex2bin(out, outLen, hex.data(), hex.size(), nullptr, &binLen, nullptr) != 0) {
        return false;
    }
    return binLen == outLen;
}

// Decode a variable-length hex string into bytes.
bool hexToBytes(const std::string& hex, std::vector<unsigned char>& out) {
    if (hex.empty() || (hex.size() % 2) != 0) return false;
    out.resize(hex.size() / 2);
    std::size_t binLen = 0;
    if (sodium_hex2bin(out.data(), out.size(), hex.data(), hex.size(), nullptr, &binLen, nullptr) != 0) {
        return false;
    }
    out.resize(binLen);
    return true;
}

std::string bytesToHex(const unsigned char* bytes, std::size_t n) {
    std::string hex(n * 2 + 1, '\0');
    sodium_bin2hex(&hex[0], hex.size(), bytes, n);
    hex.resize(n * 2);  // drop the NUL sodium_bin2hex appends
    return hex;
}

} // namespace

const char* chatCryptoStatusName(ChatCryptoStatus status) {
    switch (status) {
        case ChatCryptoStatus::Ok:                 return "Ok";
        case ChatCryptoStatus::SodiumInitFailed:   return "SodiumInitFailed";
        case ChatCryptoStatus::BadPeerKey:         return "BadPeerKey";
        case ChatCryptoStatus::BadHeaderHex:       return "BadHeaderHex";
        case ChatCryptoStatus::BadCiphertextHex:   return "BadCiphertextHex";
        case ChatCryptoStatus::CiphertextTooShort: return "CiphertextTooShort";
        case ChatCryptoStatus::SessionKeyFailed:   return "SessionKeyFailed";
        case ChatCryptoStatus::EncryptFailed:      return "EncryptFailed";
        case ChatCryptoStatus::DecryptFailed:      return "DecryptFailed";
    }
    return "Unknown";
}

void wipeChatKeyPair(ChatKeyPair& keys) {
    sodium_memzero(keys.public_key.data(), keys.public_key.size());
    sodium_memzero(keys.secret_key.data(), keys.secret_key.size());
}

ChatCryptoStatus encryptOutgoing(const ChatKeyPair& mine,
                                 const std::string& peerPublicKeyHex,
                                 const std::string& plaintext,
                                 std::string& outStreamHeaderHex,
                                 std::string& outCiphertextHex) {
    static_assert(kStreamHeaderBytes == crypto_secretstream_xchacha20poly1305_HEADERBYTES, "");
    static_assert(kStreamABytes == crypto_secretstream_xchacha20poly1305_ABYTES, "");

    if (sodium_init() < 0) return ChatCryptoStatus::SodiumInitFailed;

    unsigned char peerPk[crypto_kx_PUBLICKEYBYTES];
    if (!hexToFixed(peerPublicKeyHex, peerPk, sizeof peerPk)) return ChatCryptoStatus::BadPeerKey;

    unsigned char rx[crypto_kx_SESSIONKEYBYTES];
    unsigned char tx[crypto_kx_SESSIONKEYBYTES];
    if (crypto_kx_server_session_keys(rx, tx, mine.public_key.data(), mine.secret_key.data(), peerPk) != 0) {
        sodium_memzero(rx, sizeof rx);
        sodium_memzero(tx, sizeof tx);
        return ChatCryptoStatus::SessionKeyFailed;
    }

    crypto_secretstream_xchacha20poly1305_state state;
    unsigned char header[crypto_secretstream_xchacha20poly1305_HEADERBYTES];
    ChatCryptoStatus result = ChatCryptoStatus::EncryptFailed;

    if (crypto_secretstream_xchacha20poly1305_init_push(&state, header, tx) == 0) {
        std::vector<unsigned char> ciphertext(plaintext.size() + crypto_secretstream_xchacha20poly1305_ABYTES);
        unsigned long long ctLen = 0;
        // Only report Ok if the push actually succeeded — otherwise ctLen stays 0 and we would
        // ship a valid header with an empty ciphertext.
        if (crypto_secretstream_xchacha20poly1305_push(
                &state, ciphertext.data(), &ctLen,
                reinterpret_cast<const unsigned char*>(plaintext.data()), plaintext.size(),
                nullptr, 0, crypto_secretstream_xchacha20poly1305_TAG_FINAL) == 0) {
            outStreamHeaderHex = bytesToHex(header, sizeof header);
            outCiphertextHex = bytesToHex(ciphertext.data(), static_cast<std::size_t>(ctLen));
            result = ChatCryptoStatus::Ok;
        }
    }

    sodium_memzero(rx, sizeof rx);
    sodium_memzero(tx, sizeof tx);
    sodium_memzero(&state, sizeof state);
    return result;
}

ChatCryptoStatus decryptIncoming(const ChatKeyPair& mine,
                                 const std::string& peerPublicKeyHex,
                                 const std::string& streamHeaderHex,
                                 const std::string& ciphertextHex,
                                 std::string& outPlaintext) {
    if (sodium_init() < 0) return ChatCryptoStatus::SodiumInitFailed;

    unsigned char peerPk[crypto_kx_PUBLICKEYBYTES];
    if (!hexToFixed(peerPublicKeyHex, peerPk, sizeof peerPk)) return ChatCryptoStatus::BadPeerKey;

    unsigned char header[crypto_secretstream_xchacha20poly1305_HEADERBYTES];
    if (!hexToFixed(streamHeaderHex, header, sizeof header)) return ChatCryptoStatus::BadHeaderHex;

    std::vector<unsigned char> ciphertext;
    if (!hexToBytes(ciphertextHex, ciphertext)) return ChatCryptoStatus::BadCiphertextHex;
    // Guard the size_t subtraction below (a ciphertext shorter than the auth tag can't be
    // authentic). Exactly ABYTES is the valid empty-plaintext case, so it round-trips
    // symmetrically with encryptOutgoing (message-content policy belongs to the caller).
    if (ciphertext.size() < crypto_secretstream_xchacha20poly1305_ABYTES) {
        return ChatCryptoStatus::CiphertextTooShort;
    }

    unsigned char rx[crypto_kx_SESSIONKEYBYTES];
    unsigned char tx[crypto_kx_SESSIONKEYBYTES];
    if (crypto_kx_client_session_keys(rx, tx, mine.public_key.data(), mine.secret_key.data(), peerPk) != 0) {
        sodium_memzero(rx, sizeof rx);
        sodium_memzero(tx, sizeof tx);
        return ChatCryptoStatus::SessionKeyFailed;
    }

    crypto_secretstream_xchacha20poly1305_state state;
    ChatCryptoStatus result = ChatCryptoStatus::DecryptFailed;

    if (crypto_secretstream_xchacha20poly1305_init_pull(&state, header, rx) == 0) {
        std::vector<unsigned char> plain(ciphertext.size() - crypto_secretstream_xchacha20poly1305_ABYTES);
        unsigned long long plainLen = 0;
        unsigned char tag = 0;
        if (crypto_secretstream_xchacha20poly1305_pull(
                &state, plain.data(), &plainLen, &tag,
                ciphertext.data(), ciphertext.size(), nullptr, 0) == 0 &&
            tag == crypto_secretstream_xchacha20poly1305_TAG_FINAL) {
            outPlaintext.assign(reinterpret_cast<const char*>(plain.data()),
                                static_cast<std::size_t>(plainLen));
            result = ChatCryptoStatus::Ok;
        }
        sodium_memzero(plain.data(), plain.size());  // wipe the decrypted scratch
    }

    sodium_memzero(rx, sizeof rx);
    sodium_memzero(tx, sizeof tx);
    sodium_memzero(&state, sizeof state);
    return result;
}

} // namespace dragonx::chat