//! Public cryptographic primitives for HCFS.
//!
//! These functions implement the XChaCha20-Poly1305 streaming encryption
//! format used by HCFS. They are extracted from `Drive` so that other
//! crates (e.g. `arion-bench`) can exercise the real encryption code path
//! without constructing a full `Drive` instance.

use crate::sync::SyncError;
use blake3::Hasher as Blake3Hasher;
use chacha20poly1305::{
    XChaCha20Poly1305,
    aead::{Aead, AeadCore, AeadInPlace, KeyInit, OsRng},
};
use std::io::{Cursor, Read, Write};

/// Chunk size for streaming encryption (256KB).
pub const ENCRYPTION_CHUNK_SIZE: usize = 256 * 1024;

/// Maximum valid chunk size in ciphertext (chunk + 16-byte Poly1305 auth tag).
pub const MAX_CIPHERTEXT_CHUNK_SIZE: usize = ENCRYPTION_CHUNK_SIZE + 16;

/// Derive a per-chunk nonce by XORing the base nonce with the chunk index.
///
/// XOR only modifies the first 8 bytes; the remaining 16 bytes provide
/// file-level uniqueness from the random base nonce.
pub fn derive_chunk_nonce(base_nonce: &[u8; 24], chunk_index: u64) -> [u8; 24] {
    let mut nonce = *base_nonce;
    let index_bytes = chunk_index.to_le_bytes();
    for (i, &b) in index_bytes.iter().enumerate() {
        nonce[i] ^= b;
    }
    nonce
}

/// Read up to `buffer.len()` bytes, returning actual bytes read (0 at EOF).
/// Unlike `read_exact`, this doesn't error on partial reads at EOF.
pub(crate) fn read_exact_or_eof<R: Read>(
    reader: &mut R,
    buffer: &mut [u8],
) -> std::io::Result<usize> {
    let mut total_read = 0;
    while total_read < buffer.len() {
        match reader.read(&mut buffer[total_read..]) {
            Ok(0) => break,
            Ok(n) => total_read += n,
            Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        }
    }
    Ok(total_read)
}

/// Encrypt from reader to writer with known file size, computing ciphertext
/// hash incrementally. Streams directly without buffering all chunks in memory.
///
/// Returns the base nonce used for encryption.
///
/// Wire format:
/// ```text
/// [base_nonce: 24 bytes]
/// [chunk_count: u32 LE]
/// Per chunk:
///   [ciphertext_len: u32 LE]  (plaintext_len + 16-byte auth tag)
///   [ciphertext + auth_tag]
/// ```
pub fn encrypt_stream_with_hash<R: Read, W: Write, F>(
    reader: &mut R,
    writer: &mut W,
    key: &[u8; 32],
    file_size: u64,
    hasher: &mut Blake3Hasher,
    progress: Option<F>,
) -> Result<[u8; 24], SyncError>
where
    F: Fn(u64, u64),
{
    let cipher = XChaCha20Poly1305::new(key.into());
    let base_nonce: [u8; 24] = XChaCha20Poly1305::generate_nonce(&mut OsRng).into();

    if let Some(ref p) = progress {
        p(0, file_size);
    }

    let chunk_count = if file_size == 0 {
        1u32
    } else {
        let count = file_size.div_ceil(ENCRYPTION_CHUNK_SIZE as u64);
        if count > u32::MAX as u64 {
            return Err(SyncError::Encryption(
                "File too large: chunk count exceeds u32::MAX".to_string(),
            ));
        }
        count as u32
    };

    // Write header and update hash
    hasher.update(&base_nonce);
    writer.write_all(&base_nonce)?;

    let chunk_count_bytes = chunk_count.to_le_bytes();
    hasher.update(&chunk_count_bytes);
    writer.write_all(&chunk_count_bytes)?;

    // Handle empty file
    if file_size == 0 {
        let chunk_nonce = derive_chunk_nonce(&base_nonce, 0);
        let nonce = chacha20poly1305::XNonce::from_slice(&chunk_nonce);
        let ciphertext = cipher
            .encrypt(nonce, &[][..])
            .map_err(|e| SyncError::Encryption(format!("Empty chunk encryption failed: {}", e)))?;

        let len_bytes = (ciphertext.len() as u32).to_le_bytes();
        hasher.update(&len_bytes);
        writer.write_all(&len_bytes)?;
        hasher.update(&ciphertext);
        writer.write_all(&ciphertext)?;

        if let Some(ref p) = progress {
            p(0, 0);
        }

        return Ok(base_nonce);
    }

    // Stream encrypt chunks directly
    let mut buffer = vec![0u8; ENCRYPTION_CHUNK_SIZE];
    let mut bytes_encrypted = 0u64;

    for chunk_index in 0..chunk_count as u64 {
        let bytes_read = read_exact_or_eof(reader, &mut buffer)?;
        if bytes_read == 0 {
            return Err(SyncError::Encryption(format!(
                "Unexpected EOF at chunk {}, expected {} chunks",
                chunk_index, chunk_count
            )));
        }

        let chunk_nonce = derive_chunk_nonce(&base_nonce, chunk_index);
        let nonce = chacha20poly1305::XNonce::from_slice(&chunk_nonce);

        let tag = cipher
            .encrypt_in_place_detached(nonce, &[], &mut buffer[..bytes_read])
            .map_err(|e| {
                SyncError::Encryption(format!("Chunk {} encryption failed: {}", chunk_index, e))
            })?;

        let ciphertext_len = (bytes_read + 16) as u32;
        let len_bytes = ciphertext_len.to_le_bytes();
        hasher.update(&len_bytes);
        writer.write_all(&len_bytes)?;

        hasher.update(&buffer[..bytes_read]);
        writer.write_all(&buffer[..bytes_read])?;

        hasher.update(tag.as_slice());
        writer.write_all(tag.as_slice())?;

        bytes_encrypted += bytes_read as u64;
        if let Some(ref p) = progress {
            p(bytes_encrypted, file_size);
        }
    }

    Ok(base_nonce)
}

/// Encrypt small data (< 256KB) using chunked XChaCha20-Poly1305.
///
/// **WARNING: Buffers all chunks in memory before writing.**
/// Only use for small data like file paths.
///
/// Wire format: `[base_nonce: 24][chunk_count: u32 LE][chunks...]`
/// Each chunk: `[chunk_len: u32 LE][ciphertext + auth_tag]`
pub fn encrypt_buffered<R: Read, W: Write>(
    reader: &mut R,
    writer: &mut W,
    key: &[u8; 32],
) -> Result<[u8; 24], SyncError> {
    let cipher = XChaCha20Poly1305::new(key.into());
    let base_nonce: [u8; 24] = XChaCha20Poly1305::generate_nonce(&mut OsRng).into();

    let mut chunks: Vec<(Vec<u8>, chacha20poly1305::Tag)> = Vec::new();
    let mut buffer = vec![0u8; ENCRYPTION_CHUNK_SIZE];
    let mut chunk_index: u64 = 0;

    loop {
        let bytes_read = read_exact_or_eof(reader, &mut buffer)?;
        if bytes_read == 0 {
            break;
        }

        let chunk_nonce = derive_chunk_nonce(&base_nonce, chunk_index);
        let nonce = chacha20poly1305::XNonce::from_slice(&chunk_nonce);

        let mut ciphertext = buffer[..bytes_read].to_vec();
        let tag = cipher
            .encrypt_in_place_detached(nonce, &[], &mut ciphertext)
            .map_err(|e| {
                SyncError::Encryption(format!("Chunk {} encryption failed: {}", chunk_index, e))
            })?;

        chunks.push((ciphertext, tag));
        chunk_index += 1;
    }

    // Handle empty input
    if chunks.is_empty() {
        let chunk_nonce = derive_chunk_nonce(&base_nonce, 0);
        let nonce = chacha20poly1305::XNonce::from_slice(&chunk_nonce);
        let tag = cipher
            .encrypt_in_place_detached(nonce, &[], &mut [])
            .map_err(|e| SyncError::Encryption(format!("Empty chunk encryption failed: {}", e)))?;
        chunks.push((Vec::new(), tag));
    }

    // Write header
    writer.write_all(&base_nonce)?;
    writer.write_all(&(chunks.len() as u32).to_le_bytes())?;

    // Write chunks
    for (ciphertext, tag) in &chunks {
        let chunk_len = (ciphertext.len() + 16) as u32;
        writer.write_all(&chunk_len.to_le_bytes())?;
        writer.write_all(ciphertext)?;
        writer.write_all(tag.as_slice())?;
    }

    Ok(base_nonce)
}

/// Decrypt from reader to writer using chunked XChaCha20-Poly1305.
///
/// Expects format: `[base_nonce: 24][chunk_count: u32 LE][chunks...]`
///
/// Returns the total number of plaintext bytes written.
pub fn decrypt_stream<R: Read, W: Write, F>(
    reader: &mut R,
    writer: &mut W,
    key: &[u8; 32],
    ciphertext_size: Option<u64>,
    progress: Option<F>,
) -> Result<u64, SyncError>
where
    F: Fn(u64, u64),
{
    let cipher = XChaCha20Poly1305::new(key.into());
    let total_size = ciphertext_size.unwrap_or(0);

    if let Some(ref p) = progress {
        p(0, total_size);
    }

    // Read header
    let mut base_nonce = [0u8; 24];
    reader
        .read_exact(&mut base_nonce)
        .map_err(|e| SyncError::Decryption(format!("Failed to read nonce: {}", e)))?;

    let mut chunk_count_bytes = [0u8; 4];
    reader
        .read_exact(&mut chunk_count_bytes)
        .map_err(|e| SyncError::Decryption(format!("Failed to read chunk count: {}", e)))?;
    let chunk_count = u32::from_le_bytes(chunk_count_bytes);

    let mut bytes_read = 24u64 + 4;
    let mut total_decrypted = 0u64;

    let mut buffer = vec![0u8; MAX_CIPHERTEXT_CHUNK_SIZE];
    let mut len_bytes = [0u8; 4];

    for chunk_index in 0..chunk_count as u64 {
        reader.read_exact(&mut len_bytes).map_err(|e| {
            SyncError::Decryption(format!(
                "Failed to read chunk {} length: {}",
                chunk_index, e
            ))
        })?;
        let chunk_len = u32::from_le_bytes(len_bytes) as usize;
        bytes_read += 4;

        if chunk_len > MAX_CIPHERTEXT_CHUNK_SIZE {
            return Err(SyncError::Decryption(format!(
                "Chunk {} has invalid length {} (max {})",
                chunk_index, chunk_len, MAX_CIPHERTEXT_CHUNK_SIZE
            )));
        }

        if chunk_len < 16 {
            return Err(SyncError::Decryption(format!(
                "Chunk {} too short for auth tag: {} bytes",
                chunk_index, chunk_len
            )));
        }

        reader.read_exact(&mut buffer[..chunk_len]).map_err(|e| {
            SyncError::Decryption(format!("Failed to read chunk {} data: {}", chunk_index, e))
        })?;
        bytes_read += chunk_len as u64;

        let tag_start = chunk_len - 16;
        let mut tag_bytes = [0u8; 16];
        tag_bytes.copy_from_slice(&buffer[tag_start..chunk_len]);
        let tag = chacha20poly1305::Tag::from(tag_bytes);

        let chunk_nonce = derive_chunk_nonce(&base_nonce, chunk_index);
        let nonce = chacha20poly1305::XNonce::from_slice(&chunk_nonce);

        cipher
            .decrypt_in_place_detached(nonce, &[], &mut buffer[..tag_start], &tag)
            .map_err(|e| {
                SyncError::Decryption(format!("Chunk {} decryption failed: {}", chunk_index, e))
            })?;

        writer.write_all(&buffer[..tag_start])?;
        total_decrypted += tag_start as u64;

        if let Some(ref p) = progress {
            p(bytes_read, total_size);
        }
    }

    Ok(total_decrypted)
}

/// Encrypt small data (like paths) in memory using the buffered format.
pub fn encrypt_small(plaintext: &[u8], key: &[u8; 32]) -> Result<Vec<u8>, SyncError> {
    let mut output = Vec::new();
    encrypt_buffered(&mut Cursor::new(plaintext), &mut output, key)?;
    Ok(output)
}

/// Decrypt small data (like paths) in memory using the streaming format.
pub fn decrypt_small(ciphertext: &[u8], key: &[u8; 32]) -> Result<Vec<u8>, SyncError> {
    let mut output = Vec::new();
    decrypt_stream(
        &mut Cursor::new(ciphertext),
        &mut output,
        key,
        None,
        None::<fn(u64, u64)>,
    )?;
    Ok(output)
}

/// Compute BLAKE3 of a path string, returning a 32-byte path hash.
pub fn compute_path_hash(path: &str) -> [u8; 32] {
    blake3::hash(path.as_bytes()).into()
}

/// Compute salted content hash using streaming (memory-efficient).
///
/// Returns `(salted_hash, total_bytes_read)`.
/// The hash is `BLAKE3(user_id || file_content)`.
pub fn compute_salted_hash<R: Read>(
    reader: &mut R,
    user_id: &str,
) -> Result<([u8; 32], u64), SyncError> {
    let mut hasher = Blake3Hasher::new();
    hasher.update(user_id.as_bytes());

    let mut buffer = vec![0u8; ENCRYPTION_CHUNK_SIZE];
    let mut total_size: u64 = 0;

    loop {
        let bytes_read = reader.read(&mut buffer)?;
        if bytes_read == 0 {
            break;
        }
        hasher.update(&buffer[..bytes_read]);
        total_size += bytes_read as u64;
    }

    let salted_hash: [u8; 32] = hasher.finalize().into();
    Ok((salted_hash, total_size))
}
