//! Encryption — encrypt, decrypt, permissions, and encrypted-status queries.
//!
//! Wraps the low-level primitives in `crate::enhanced::encryption` behind the
//! cpdf `CpdfDocument` interface.  Zero `unsafe` in this file.

use crate::cpdf::document::CpdfDocument;
use crate::cpdf::error::{CpdfError, Result};
use crate::enhanced::encryption::{
    EncryptionAlgorithm, EncryptionInfo, EncryptionOptions, PdfDecryptor, PdfEncryptor,
    Permissions, get_encryption_info_from_data,
};

// ── Re-exports for convenience ──────────────────────────────────────

/// Encryption method mirroring `enhanced::EncryptionAlgorithm` but living in
/// cpdf's public surface.  A thin newtype so that callers don't need to import
/// the enhanced module.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EncryptionMethod {
    Rc4_40,
    Rc4_128,
    Aes128,
    Aes256,
}

impl EncryptionMethod {
    fn to_algorithm(self) -> EncryptionAlgorithm {
        match self {
            EncryptionMethod::Rc4_40 => EncryptionAlgorithm::Rc4_40,
            EncryptionMethod::Rc4_128 => EncryptionAlgorithm::Rc4_128,
            EncryptionMethod::Aes128 => EncryptionAlgorithm::Aes128,
            EncryptionMethod::Aes256 => EncryptionAlgorithm::Aes256,
        }
    }

    pub fn from_code(code: i32) -> Result<Self> {
        match code {
            0 => Ok(EncryptionMethod::Rc4_40),
            1 => Ok(EncryptionMethod::Rc4_128),
            2 => Ok(EncryptionMethod::Aes128),
            3 => Ok(EncryptionMethod::Aes256),
            _ => Err(CpdfError::InvalidArgument(format!(
                "unknown encryption method code: {code}"
            ))),
        }
    }
}

// ── Permission bits ─────────────────────────────────────────────────

/// Permission flags as a simple u32 bitfield.  Matches the PDF spec.
pub const PERM_PRINT: u32 = 1 << 2;
pub const PERM_MODIFY: u32 = 1 << 3;
pub const PERM_COPY: u32 = 1 << 4;
pub const PERM_ANNOTATE: u32 = 1 << 5;
pub const PERM_FORMS: u32 = 1 << 8;
pub const PERM_ACCESSIBILITY: u32 = 1 << 9;
pub const PERM_ASSEMBLE: u32 = 1 << 10;
pub const PERM_PRINT_HIGH: u32 = 1 << 11;

fn bits_to_permissions(bits: u32) -> Permissions {
    Permissions {
        print: bits & PERM_PRINT != 0,
        modify: bits & PERM_MODIFY != 0,
        copy: bits & PERM_COPY != 0,
        annotate: bits & PERM_ANNOTATE != 0,
        fill_forms: bits & PERM_FORMS != 0,
        extract_accessibility: bits & PERM_ACCESSIBILITY != 0,
        assemble: bits & PERM_ASSEMBLE != 0,
        print_high_quality: bits & PERM_PRINT_HIGH != 0,
    }
}

fn permissions_to_bits(p: &Permissions) -> u32 {
    let mut bits: u32 = 0;
    if p.print {
        bits |= PERM_PRINT;
    }
    if p.modify {
        bits |= PERM_MODIFY;
    }
    if p.copy {
        bits |= PERM_COPY;
    }
    if p.annotate {
        bits |= PERM_ANNOTATE;
    }
    if p.fill_forms {
        bits |= PERM_FORMS;
    }
    if p.extract_accessibility {
        bits |= PERM_ACCESSIBILITY;
    }
    if p.assemble {
        bits |= PERM_ASSEMBLE;
    }
    if p.print_high_quality {
        bits |= PERM_PRINT_HIGH;
    }
    bits
}

// ── Public API ──────────────────────────────────────────────────────

/// Encrypt a document with the given method, passwords, and permission bits.
pub fn encrypt(
    doc: CpdfDocument,
    method: EncryptionMethod,
    user_pw: &str,
    owner_pw: &str,
    perms_bits: u32,
) -> Result<CpdfDocument> {
    let perms = bits_to_permissions(perms_bits);
    let opts = EncryptionOptions::new()
        .algorithm(method.to_algorithm())
        .user_password(user_pw)
        .owner_password(owner_pw)
        .permissions(perms);

    let mut encryptor = PdfEncryptor::new(doc.to_bytes(), opts)
        .map_err(|e| CpdfError::Encryption(e.to_string()))?;

    let encrypted = encryptor
        .encrypt()
        .map_err(|e| CpdfError::Encryption(e.to_string()))?;

    CpdfDocument::from_bytes(encrypted)
}

/// Decrypt a document with the given password.
pub fn decrypt(doc: &CpdfDocument, password: &str) -> Result<CpdfDocument> {
    let mut decryptor =
        PdfDecryptor::new(doc.to_bytes()).map_err(|e| CpdfError::Encryption(e.to_string()))?;

    let user_ok = decryptor
        .authenticate_user(password)
        .map_err(|e| CpdfError::Encryption(e.to_string()))?;
    if !user_ok {
        let owner_ok = decryptor
            .authenticate_owner(password)
            .map_err(|e| CpdfError::Encryption(e.to_string()))?;
        if !owner_ok {
            return Err(CpdfError::Encryption("invalid password".into()));
        }
    }

    let decrypted = decryptor
        .decrypt()
        .map_err(|e| CpdfError::Encryption(e.to_string()))?;

    CpdfDocument::from_bytes(decrypted)
}

/// Returns `true` if the document has an `/Encrypt` dictionary.
pub fn is_encrypted(doc: &CpdfDocument) -> bool {
    let bytes = doc.to_bytes();
    let text = String::from_utf8_lossy(&bytes);
    text.contains("/Encrypt")
}

/// Returns the `EncryptionInfo` for the document.
pub fn encryption_info(doc: &CpdfDocument) -> Result<EncryptionInfo> {
    get_encryption_info_from_data(&doc.to_bytes()).map_err(|e| CpdfError::Encryption(e.to_string()))
}

/// Check whether a specific permission bit is granted according to the
/// document's `/Encrypt` dictionary.  Returns `true` if the document is
/// unencrypted (all permissions implicitly granted).
pub fn has_permission(doc: &CpdfDocument, _password: &str, perm_bit: u32) -> Result<bool> {
    let info = get_encryption_info_from_data(&doc.to_bytes())
        .map_err(|e| CpdfError::Encryption(e.to_string()))?;

    if !info.is_encrypted {
        return Ok(true);
    }

    match info.permissions {
        Some(p) => {
            let bits = permissions_to_bits(&p);
            Ok(bits & perm_bit != 0)
        }
        None => Ok(false),
    }
}

/// Convenience: decrypt a document if it is encrypted, pass through unchanged
/// if it is not.
pub fn decrypt_if_needed(doc: CpdfDocument, password: &str) -> Result<CpdfDocument> {
    if is_encrypted(&doc) {
        decrypt(&doc, password)
    } else {
        Ok(doc)
    }
}

// ── Tests ───────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    fn blank() -> CpdfDocument {
        CpdfDocument::blank(612.0, 792.0, 1).unwrap()
    }

    #[test]
    fn unencrypted_doc_reports_false() {
        let doc = blank();
        assert!(!is_encrypted(&doc));
    }

    #[test]
    fn unencrypted_has_all_permissions() {
        let doc = blank();
        assert!(has_permission(&doc, "", PERM_PRINT).unwrap());
        assert!(has_permission(&doc, "", PERM_COPY).unwrap());
    }

    #[test]
    fn decrypt_if_needed_passthrough() {
        let doc = blank();
        let out = decrypt_if_needed(doc.clone(), "").unwrap();
        assert_eq!(doc.page_count(), out.page_count());
    }

    #[test]
    fn encryption_method_from_code() {
        assert_eq!(
            EncryptionMethod::from_code(0).unwrap(),
            EncryptionMethod::Rc4_40
        );
        assert_eq!(
            EncryptionMethod::from_code(3).unwrap(),
            EncryptionMethod::Aes256
        );
        assert!(EncryptionMethod::from_code(99).is_err());
    }

    #[test]
    fn permissions_roundtrip() {
        let perms = bits_to_permissions(PERM_PRINT | PERM_COPY | PERM_FORMS);
        let bits = permissions_to_bits(&perms);
        assert_eq!(bits, PERM_PRINT | PERM_COPY | PERM_FORMS);
    }

    #[test]
    fn encrypt_produces_encrypt_dict() {
        let doc = blank();
        let enc = encrypt(
            doc,
            EncryptionMethod::Aes256,
            "user",
            "owner",
            PERM_PRINT | PERM_COPY,
        )
        .unwrap();
        assert!(is_encrypted(&enc));
    }
}
