import sodium, { type StateAddress } from "libsodium-wrappers-sumo";

// Plaintext chunk size used by Ente for file content streams. Hard-coded by
// the server; clients must match.
export const STREAM_CHUNK_SIZE = 4 * 1024 * 1024;

// Per-chunk overhead added by libsodium's secretstream construction:
// 16 bytes of Poly1305 tag plus 1 byte of secretstream tag.
export const STREAM_CHUNK_OVERHEAD = 17;

// libsodium's crypto_secretstream_xchacha20poly1305_TAG_FINAL: the tag that
// marks the last chunk of a stream. Exported so callers (the download layer)
// can detect truncation without importing sodium themselves.
//
// This is a function rather than a constant, and that is load-bearing:
// libsodium attaches its constants to the module object inside
// `ready.then(...)`, which resolves long after this module is evaluated. A
// module-level read would bind `undefined`, every downstream comparison
// against it would then be false, and every valid download would be rejected
// as truncated. Reading at call time returns the library's own value, so
// there is also no second copy of a protocol constant to keep in sync.
//
// Under vitest sodium is already initialised in the worker before this module
// is evaluated, so an eager read would pick up a real value there and the
// ordinary tests could not tell the difference. The regression guard is
// "streamTagFinal() reads the constant at call time, not at import time" in
// test/crypto/stream.test.ts, which reproduces the plain-Node ESM ordering
// against a stand-in sodium module; it goes red if this becomes eager.
export const streamTagFinal = (): number =>
    sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL;

// Encrypt a small blob as a single secretstream chunk with TAG_FINAL.
// Returns the header and ciphertext. Used for encrypting thumbnails
// and metadata before upload.
export const encryptBlob = (
    plaintext: Uint8Array,
    key: Uint8Array,
): { header: Uint8Array; ciphertext: Uint8Array } => {
    const push = sodium.crypto_secretstream_xchacha20poly1305_init_push(key);
    const ciphertext = sodium.crypto_secretstream_xchacha20poly1305_push(
        push.state,
        plaintext,
        null,
        sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL,
    );
    return { header: push.header, ciphertext };
};

// Opaque handle to libsodium's secretstream pull state. Threaded through
// successive pullStreamChunk calls. The type comes from the named export;
// the default import is the module's value side and has no type namespace
// under it.
export type StreamPullState = StateAddress;

// Initialise a pull stream from the per-file decryption header and the
// per-file key.
export const initStreamPull = (
    header: Uint8Array,
    key: Uint8Array,
): StreamPullState =>
    sodium.crypto_secretstream_xchacha20poly1305_init_pull(header, key);

// Decrypt a small blob that was encrypted as a single secretstream chunk
// with TAG_FINAL. Ente uses this form ("blob") for file metadata and
// magic metadata — anything under ~1 MiB that isn't chunked.
export const decryptBlob = (
    ciphertext: Uint8Array,
    header: Uint8Array,
    key: Uint8Array,
): Uint8Array => {
    const state = initStreamPull(header, key);
    const { plaintext, tag } = pullStreamChunk(state, ciphertext);
    const tagFinal = streamTagFinal();
    if (tag !== tagFinal) {
        throw new Error(
            `decryptBlob: expected TAG_FINAL (${tagFinal}), got tag ${tag}`,
        );
    }
    return plaintext;
};

// Decrypt one ciphertext chunk. Returns the plaintext and the secretstream
// tag (0=MESSAGE, 1=PUSH, 2=REKEY, 3=FINAL). The caller must verify the
// stream ended on TAG_FINAL to detect truncation.
export const pullStreamChunk = (
    state: StreamPullState,
    ciphertext: Uint8Array,
): { plaintext: Uint8Array; tag: number } => {
    // The additional-data argument is not optional in libsodium's signature.
    // null is "no additional data", matching the null passed on the push side
    // in encryptBlob; Ente's file streams carry none.
    const result = sodium.crypto_secretstream_xchacha20poly1305_pull(
        state,
        ciphertext,
        null,
    );
    if (result === false) {
        throw new Error("secretstream chunk authentication failed");
    }
    return { plaintext: result.message, tag: result.tag };
};