/*
 * dm_crypto.c — криптография прямого чата (DM)
 */

#include "dm_crypto.h"
#include "../../lib/debug_config.h"

#define OPENSSL_API_COMPAT 0x10100000L
#include <openssl/evp.h>
#include <openssl/sha.h>

#include <string.h>
#include <stdio.h>

#define DM_ID "dm_crypto"

/* Медиа использует тот же ключ target, но отдельное пространство nonce. */
void dm_build_media_nonce(uint64_t author, const uint8_t media_id[16], uint64_t part, uint8_t nonce[DM_NONCE_SIZE]) {
    SHA256_CTX ctx;
    uint8_t hash[32];
    SHA256_Init(&ctx);
    SHA256_Update(&ctx, "utun_dm_media", 13);
    SHA256_Update(&ctx, &author, sizeof(author));
    SHA256_Update(&ctx, media_id, 16);
    SHA256_Update(&ctx, &part, sizeof(part));
    SHA256_Final(hash, &ctx);
    memcpy(nonce, hash, DM_NONCE_SIZE);
}

uint64_t dm_derive_conv_id(uint64_t a, uint64_t b) {
    uint64_t lo = a < b ? a : b;
    uint64_t hi = a < b ? b : a;

    SHA256_CTX ctx;
    uint8_t hash[32];
    SHA256_Init(&ctx);
    SHA256_Update(&ctx, "utun_dm_v1", 9);
    SHA256_Update(&ctx, &lo, 8);
    SHA256_Update(&ctx, &hi, 8);
    SHA256_Final(hash, &ctx);

    uint64_t id;
    memcpy(&id, hash, 8);
    id &= 0x7FFFFFFFFFFFFFFFULL;
    if (id == 0) id = 1;
    return id;
}

int dm_derive_content_key(const uint8_t my_priv[32], const uint8_t peer_pub[32],
                          uint8_t out[DM_CONTENT_KEY_SIZE]) {
    if (!my_priv || !peer_pub || !out) return -1;

    EVP_PKEY* my_pkey = EVP_PKEY_new_raw_private_key(EVP_PKEY_X25519, NULL, my_priv, 32);
    if (!my_pkey) { DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: EVP_PKEY_new_raw_private_key failed", DM_ID); return -1; }
    EVP_PKEY* peer_pkey = EVP_PKEY_new_raw_public_key(EVP_PKEY_X25519, NULL, peer_pub, 32);
    if (!peer_pkey) { EVP_PKEY_free(my_pkey); DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: EVP_PKEY_new_raw_public_key failed", DM_ID); return -1; }

    EVP_PKEY_CTX* dctx = EVP_PKEY_CTX_new(my_pkey, NULL);
    uint8_t shared[32];
    size_t shared_len = sizeof(shared);
    int rc = -1;
    if (!dctx) goto out;
    if (EVP_PKEY_derive_init(dctx) <= 0) goto out;
    if (EVP_PKEY_derive_set_peer(dctx, peer_pkey) <= 0) goto out;
    if (EVP_PKEY_derive(dctx, shared, &shared_len) <= 0 || shared_len != sizeof(shared)) goto out;

    SHA256_CTX sctx;
    SHA256_Init(&sctx);
    SHA256_Update(&sctx, shared, sizeof(shared));
    SHA256_Update(&sctx, "utun_dm_content", 15);
    SHA256_Final(out, &sctx);
    rc = 0;

out:
    if (dctx) EVP_PKEY_CTX_free(dctx);
    EVP_PKEY_free(my_pkey);
    EVP_PKEY_free(peer_pkey);
    if (rc != 0) DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: X25519 derive failed", DM_ID);
    return rc;
}

void dm_build_nonce(uint64_t conv_id, uint64_t author, uint64_t seq, uint8_t nonce[DM_NONCE_SIZE]) {
    SHA256_CTX ctx;
    uint8_t hash[32];
    SHA256_Init(&ctx);
    SHA256_Update(&ctx, &conv_id, 8);
    SHA256_Update(&ctx, &author, 8);
    SHA256_Update(&ctx, &seq, 8);
    SHA256_Final(hash, &ctx);
    memcpy(nonce, hash, DM_NONCE_SIZE);
}

static int dm_ccm_crypt(const uint8_t key[DM_CONTENT_KEY_SIZE], const uint8_t nonce[DM_NONCE_SIZE],
                        const uint8_t* in, size_t in_len, uint8_t* out, size_t* out_len,
                        int encrypt, const uint8_t* tag_in) {
    if (!key || !nonce || (!in && in_len) || !out || !out_len ||
        (encrypt ? in_len > DM_CCM_PLAIN_MAX : in_len < DM_TAG_SIZE || in_len - DM_TAG_SIZE > DM_CCM_PLAIN_MAX)) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: invalid CCM arguments encrypt=%d bytes=%zu", DM_ID, encrypt, in_len);
        return -1;
    }
    *out_len = 0;
    static const uint8_t empty = 0;
    if (!in) in = &empty;

    EVP_CIPHER_CTX* ectx = EVP_CIPHER_CTX_new();
    if (!ectx) { DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: CCM context allocation failed", DM_ID); return -1; }
    const EVP_CIPHER* cipher = EVP_aes_256_ccm();

    int rc = -1;
    if (encrypt) {
        if (EVP_EncryptInit_ex(ectx, cipher, NULL, NULL, NULL) != 1
            || EVP_CIPHER_CTX_ctrl(ectx, EVP_CTRL_AEAD_SET_IVLEN, DM_NONCE_SIZE, NULL) != 1
            || EVP_CIPHER_CTX_ctrl(ectx, EVP_CTRL_AEAD_SET_TAG, DM_TAG_SIZE, NULL) != 1
            || EVP_EncryptInit_ex(ectx, NULL, NULL, key, nonce) != 1) goto done;
        int outl;
        if (EVP_EncryptUpdate(ectx, NULL, &outl, NULL, (int)in_len) != 1) goto done;
        if (EVP_EncryptUpdate(ectx, out, &outl, in, (int)in_len) != 1 || outl != (int)in_len) goto done;
        int tmpl;
        if (EVP_EncryptFinal_ex(ectx, out + outl, &tmpl) != 1) goto done;
        if (EVP_CIPHER_CTX_ctrl(ectx, EVP_CTRL_AEAD_GET_TAG, DM_TAG_SIZE, out + in_len) != 1) goto done;
        *out_len = in_len + DM_TAG_SIZE;
    } else {
        if (in_len < DM_TAG_SIZE) goto done;
        size_t body = in_len - DM_TAG_SIZE;
        if (EVP_DecryptInit_ex(ectx, cipher, NULL, NULL, NULL) != 1
            || EVP_CIPHER_CTX_ctrl(ectx, EVP_CTRL_AEAD_SET_IVLEN, DM_NONCE_SIZE, NULL) != 1
            || EVP_CIPHER_CTX_ctrl(ectx, EVP_CTRL_AEAD_SET_TAG, DM_TAG_SIZE, (void*)(tag_in ? tag_in : in + body)) != 1
            || EVP_DecryptInit_ex(ectx, NULL, NULL, key, nonce) != 1) goto done;
        int outl;
        if (EVP_DecryptUpdate(ectx, NULL, &outl, NULL, (int)body) != 1) goto done;
        if (EVP_DecryptUpdate(ectx, out, &outl, in, (int)body) != 1 || outl != (int)body) goto done;
        *out_len = body;
    }
    rc = 0;
done:
    EVP_CIPHER_CTX_free(ectx);
    if (rc) DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: CCM %s failed bytes=%zu", DM_ID, encrypt ? "encrypt" : "decrypt", in_len);
    return rc;
}

int dm_encrypt(const uint8_t key[DM_CONTENT_KEY_SIZE], const uint8_t nonce[DM_NONCE_SIZE],
               const uint8_t* plain, size_t plain_len, uint8_t* out, size_t* out_len) {
    return dm_ccm_crypt(key, nonce, plain, plain_len, out, out_len, 1, NULL);
}

int dm_decrypt(const uint8_t key[DM_CONTENT_KEY_SIZE], const uint8_t nonce[DM_NONCE_SIZE],
               const uint8_t* ct, size_t ct_len, uint8_t* out, size_t* out_len) {
    return dm_ccm_crypt(key, nonce, ct, ct_len, out, out_len, 0, NULL);
}

/* Каждая порция имеет собственный nonce, ciphertext неизменен для всех сетевых повторов. */
int dm_media_encrypt_stream(const uint8_t key[DM_CONTENT_KEY_SIZE], uint64_t author, const uint8_t media_id[16],
                            FILE* src, FILE* dst, uint64_t* plain_size, uint8_t ciphertext_hash[32]) {
    if (!key || !author || !media_id || !src || !dst || src == dst || !plain_size || !ciphertext_hash) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: invalid media encryption arguments", DM_ID);
        return -1;
    }
    *plain_size = 0;
    memset(ciphertext_hash, 0, 32);
    uint8_t plain[DM_MEDIA_PLAIN_SIZE], cipher[DM_MEDIA_PLAIN_SIZE + DM_TAG_SIZE], nonce[DM_NONCE_SIZE];
    uint64_t part = 0, total = 0;
    SHA256_CTX hash;
    SHA256_Init(&hash);
    for (;;) {
        size_t n = fread(plain, 1, sizeof(plain), src);
        if (ferror(src)) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media source read failed part=%llu", DM_ID, (unsigned long long)part);
            return -1;
        }
        if (!n && part) break;
        size_t clen = 0;
        dm_build_media_nonce(author, media_id, part, nonce);
        if (dm_encrypt(key, nonce, plain, n, cipher, &clen) != 0) return -1;
        if (fwrite(cipher, 1, clen, dst) != clen) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media ciphertext write failed part=%llu", DM_ID, (unsigned long long)part);
            return -1;
        }
        SHA256_Update(&hash, cipher, clen);
        total += n;
        part++;
        if (n < sizeof(plain)) break;
    }
    if (fflush(dst) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media ciphertext flush failed", DM_ID);
        return -1;
    }
    SHA256_Final(ciphertext_hash, &hash);
    *plain_size = total;
    DEBUG_DEBUG(DEBUG_CATEGORY_DM, "%s: media encrypted author=%llu plain=%llu records=%llu", DM_ID,
                (unsigned long long)author, (unsigned long long)total, (unsigned long long)part);
    return 0;
}

/* Размер и хеш берутся из подписанного manifest; никакого size-only принятия файла. */
int dm_media_decrypt_stream(const uint8_t key[DM_CONTENT_KEY_SIZE], uint64_t author, const uint8_t media_id[16],
                            uint64_t plain_size, const uint8_t ciphertext_hash[32], FILE* src, FILE* dst) {
    if (!key || !author || !media_id || !ciphertext_hash || !src || !dst || src == dst) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: invalid media decryption arguments", DM_ID);
        return -1;
    }
    uint8_t plain[DM_MEDIA_PLAIN_SIZE], cipher[DM_MEDIA_PLAIN_SIZE + DM_TAG_SIZE], nonce[DM_NONCE_SIZE];
    uint64_t remaining = plain_size, part = 0;
    SHA256_CTX hash;
    SHA256_Init(&hash);
    do {
        size_t n = remaining > DM_MEDIA_PLAIN_SIZE ? DM_MEDIA_PLAIN_SIZE : (size_t)remaining;
        size_t clen = n + DM_TAG_SIZE, plen = 0;
        if (fread(cipher, 1, clen, src) != clen) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media ciphertext short/read failure part=%llu expected=%zu", DM_ID,
                        (unsigned long long)part, clen);
            return -1;
        }
        SHA256_Update(&hash, cipher, clen);
        dm_build_media_nonce(author, media_id, part, nonce);
        if (dm_decrypt(key, nonce, cipher, clen, plain, &plen) != 0) return -1;
        if (plen != n || fwrite(plain, 1, plen, dst) != plen) {
            DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media plaintext write/length failure part=%llu", DM_ID, (unsigned long long)part);
            return -1;
        }
        remaining -= n;
        part++;
    } while (remaining);
    uint8_t actual_hash[32];
    SHA256_Final(actual_hash, &hash);
    if (fgetc(src) != EOF || ferror(src) || memcmp(actual_hash, ciphertext_hash, 32)) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media ciphertext size/hash mismatch plain=%llu", DM_ID, (unsigned long long)plain_size);
        return -1;
    }
    if (fflush(dst) != 0) {
        DEBUG_ERROR(DEBUG_CATEGORY_DM, "%s: media plaintext flush failed", DM_ID);
        return -1;
    }
    DEBUG_DEBUG(DEBUG_CATEGORY_DM, "%s: media verified author=%llu plain=%llu records=%llu", DM_ID,
                (unsigned long long)author, (unsigned long long)plain_size, (unsigned long long)part);
    return 0;
}