use pdfmake_rust::pdf_writer::PdfWriter;
use pdfmake_rust::pdf_writer::{compute_encryption_key, compute_user_password_with_id, rc4_encrypt_stream, encrypt_pdf_string};
use pdfmake_rust::{DocumentInfo, PdfEncryption};

#[test]
fn test_generate_file_id() {
    let writer = PdfWriter::new();
    let id = writer.generate_file_id();
    assert_eq!(id.len(), 16);
}

#[test]
fn test_compute_encryption_key() {
    let user_pwd = "user123";
    let o_value = vec![0u8; 32]; // 示例 O 值
    let permissions = -3904i32;
    let file_id = vec![1u8; 16]; // 示例文件 ID

    let key = compute_encryption_key(user_pwd, &o_value, permissions, &file_id);
    // 加密密钥应该是 16 字节（128-bit）
    assert_eq!(key.len(), 16);
}

#[test]
fn test_compute_user_password_with_id() {
    let user_pwd = "user123";
    let o_value = vec![0u8; 32];
    let permissions = -3904i32;
    let file_id = vec![1u8; 16];

    let u_value = compute_user_password_with_id(user_pwd, &o_value, permissions, &file_id);
    // R=3 U 值应该是 32 字节（16 字节 hash + 16 字节 padding）
    assert_eq!(u_value.len(), 32);
}

#[test]
fn test_compute_user_password_with_id_matches_pdf_r3_shape() {
    let user_pwd = "123123";
    let o_value = vec![0u8; 32];
    let permissions = -3904i32;
    let file_id = vec![1u8; 16];

    let u_value = compute_user_password_with_id(user_pwd, &o_value, permissions, &file_id);
    assert_eq!(u_value.len(), 32, "R=3 U value must be 32 bytes");
    assert_eq!(&u_value[16..], &[
        0x28, 0xBF, 0x4E, 0x5E, 0x4E, 0x75, 0x8A, 0x41,
        0x64, 0x00, 0x4E, 0x56, 0xFF, 0xFA, 0x01, 0x08,
    ], "R=3 U value tail must match the standard padding prefix");
}

#[test]
fn test_compute_encryption_key_different_inputs() {
    // 测试不同输入产生不同结果
    let o_value = vec![0u8; 32];
    let permissions = -3904i32;
    let file_id = vec![1u8; 16];
    
    let key1 = compute_encryption_key("password1", &o_value, permissions, &file_id);
    let key2 = compute_encryption_key("password2", &o_value, permissions, &file_id);
    
    assert_ne!(key1, key2, "不同密码应该产生不同的加密密钥");
}

#[test]
fn test_compute_user_password_with_id_different_inputs() {
    // 测试不同输入产生不同结果
    let o_value = vec![0u8; 32];
    let permissions = -3904i32;
    let file_id = vec![1u8; 16];
    
    let u1 = compute_user_password_with_id("password1", &o_value, permissions, &file_id);
    let u2 = compute_user_password_with_id("password2", &o_value, permissions, &file_id);
    
    assert_ne!(u1, u2, "不同密码应该产生不同的U值");
}

#[test]
fn test_pdf_writer_encryption_context() {
    let mut writer = PdfWriter::new();
    
    // 初始状态：没有加密上下文
    assert!(writer.encryption_key().is_none());
    assert!(writer.file_id().is_none());
    
    // 设置加密上下文
    let key = vec![1u8, 2, 3, 4, 5]; // 5 字节（40-bit）密钥
    let file_id = vec![10u8; 16]; // 16 字节文件 ID
    writer.set_encryption_context(key.clone(), file_id.clone());
    
    // 验证可以获取加密上下文
    assert_eq!(writer.encryption_key(), Some(&key[..]));
    assert_eq!(writer.file_id(), Some(&file_id[..]));
}

#[test]
fn test_pdf_writer_encryption_context_default() {
    let writer = PdfWriter::new();
    
    // 默认 PdfWriter 应该没有加密上下文
    assert!(writer.encryption_key().is_none());
    assert!(writer.file_id().is_none());
}

#[test]
fn test_encrypt_stream_data() {
    let key = vec![0x01, 0x02, 0x03, 0x04, 0x05]; // 40-bit key
    let data = b"Hello, World! This is a test stream.";
    
    let encrypted = rc4_encrypt_stream(&key, data);
    // 加密后数据长度应与原始数据相同
    assert_eq!(encrypted.len(), data.len());
    // 加密后数据应与原始数据不同
    assert_ne!(encrypted, data);
    
    // RC4 是对称的，再次加密应恢复原始数据
    let decrypted = rc4_encrypt_stream(&key, &encrypted);
    assert_eq!(decrypted, data);
}

#[test]
fn test_encrypt_stream_data_empty_key() {
    let key = vec![]; // 空密钥
    let data = b"Hello, World!";
    
    // 空密钥应返回原始数据
    let result = rc4_encrypt_stream(&key, data);
    assert_eq!(result, data);
}

#[test]
fn test_encrypt_pdf_string() {
    let key = vec![0x01, 0x02, 0x03, 0x04, 0x05]; // 40-bit key
    let text = "Hello, World!";
    
    let encrypted = encrypt_pdf_string(&key, text);
    // 加密后应该是有效的 PDF 字符串格式
    assert!(encrypted.starts_with('('));
    assert!(encrypted.ends_with(')'));
    // 加密后内容应与原始不同
    assert!(!encrypted.contains(text));
}

#[test]
fn test_encrypt_pdf_string_empty_key() {
    let key = vec![];
    let text = "Hello, World!";
    
    let result = encrypt_pdf_string(&key, text);
    // 空密钥时应返回未加密的 PDF 字符串（包含原始文本）
    assert!(result.starts_with('('));
    assert!(result.ends_with(')'));
    assert!(result.contains(text));
}

#[test]
fn test_setup_encryption_stores_consistent_key_and_file_id() {
    let mut writer = PdfWriter::new();
    let enc = PdfEncryption {
        user_password: Some("user".to_string()),
        owner_password: Some("owner".to_string()),
        allow_print: true,
        allow_copy: false,
        allow_modify: false,
    };

    writer.setup_encryption(&enc);

    // setup_encryption should have set both key and file_id
    let key = writer.encryption_key().expect("encryption key should be set").to_vec();
    let file_id = writer.file_id().expect("file_id should be set").to_vec();
    assert_eq!(key.len(), 16, "encryption key should be 128-bit (16 bytes)");
    assert_eq!(file_id.len(), 16, "file_id should be 16 bytes");

    // Add a dummy object so finish() has something to write
    writer.add_object(b"<< /Type /Catalog >>".to_vec());

    // finish() with the same encryption config should NOT regenerate file_id
    let pdf = writer.finish(1, None, Some(&enc), None);
    let pdf_str = String::from_utf8_lossy(&pdf);

    // Extract the /ID from the trailer — it should use the same file_id
    let id_hex: String = file_id.iter().map(|b| format!("{:02X}", b)).collect();
    assert!(
        pdf_str.contains(&format!("/ID [<{}> <{}>]", id_hex, id_hex)),
        "PDF trailer /ID should use the file_id from setup_encryption, not a regenerated one"
    );
}

#[test]
fn test_setup_encryption_key_used_for_stream_encryption() {
    let mut writer = PdfWriter::new();
    let enc = PdfEncryption {
        user_password: Some("user".to_string()),
        owner_password: Some("owner".to_string()),
        allow_print: true,
        allow_copy: false,
        allow_modify: false,
    };

    writer.setup_encryption(&enc);
    let key = writer.encryption_key().expect("key should be set").to_vec();

    // Add a stream object — it should be encrypted with a per-object key
    writer.add_object_with_stream(
        b"<< /Length 5 >>".to_vec(),
        b"hello".to_vec(),
    );

    // Per-object key: MD5(global_key + object_number_le3 + gen_le2)[..10]
    // Object number is 1 (first object), generation 0
    let mut obj_key_data = Vec::new();
    obj_key_data.extend_from_slice(&key);
    obj_key_data.extend_from_slice(&1u32.to_le_bytes()[..3]);
    obj_key_data.push(0u8);
    obj_key_data.push(0u8);
    use md5::{Md5, Digest};
    let mut hasher = Md5::new();
    hasher.update(&obj_key_data);
    let obj_key = hasher.finalize()[..16].to_vec();

    let expected_encrypted = rc4_encrypt_stream(&obj_key, b"hello");
    let pdf = writer.finish(1, None, Some(&enc), None);
    // The stream should contain the encrypted bytes, not plaintext "hello"
    assert!(
        pdf.windows(expected_encrypted.len()).any(|w| w == expected_encrypted),
        "Stream should be encrypted with per-object key derived from setup_encryption key"
    );
}

#[test]
fn test_finish_generates_file_id_when_no_encryption_context() {
    // When no setup_encryption is called and no encryption passed to finish,
    // finish should generate a file_id for the trailer
    let mut writer = PdfWriter::new();
    writer.add_object(b"<< /Type /Catalog >>".to_vec());
    let pdf = writer.finish(1, None, None, None);
    let pdf_str = String::from_utf8_lossy(&pdf);
    // Should still have /ID in trailer
    assert!(pdf_str.contains("/ID ["), "Trailer should contain /ID even without encryption");
}

// ===== Integration Tests =====

use pdfmake_rust::{Document, DocumentNode, Margins, PageSize, PdfMake, Style, TextNode};

#[test]
fn test_encrypted_pdf_generation() {
    let doc = Document::builder()
        .info(DocumentInfo {
            title: Some("加密测试文档".into()),
            author: Some("测试作者".into()),
            ..Default::default()
        })
        .encryption(PdfEncryption {
            user_password: Some("user123".into()),
            owner_password: Some("owner456".into()),
            allow_print: true,
            allow_copy: false,
            allow_modify: false,
        })
        .page_size(PageSize::a4())
        .page_margins(Margins::new(50.0, 50.0, 50.0, 50.0))
        .content(DocumentNode::Text(TextNode::new("加密测试内容").inline_style(Style::new().font_size(14.0))))
        .build();

    let pdf_bytes = PdfMake::new().render(&doc).expect("生成 PDF 失败");

    // 验证 PDF 基本结构
    assert!(pdf_bytes.len() > 100);
    let text = String::from_utf8_lossy(&pdf_bytes);

    // 验证加密字典存在
    assert!(text.contains("/Encrypt"));
    assert!(text.contains("/Standard"));
    assert!(text.contains("/V 2"));
    assert!(text.contains("/R 3"));
    assert!(text.contains("/Length 128"));

    // 验证文件 ID 存在
    assert!(text.contains("/ID"));

    // 保存文件用于手动验证
    std::fs::write("tests/output/encrypted_test.pdf", &pdf_bytes).unwrap();
}

#[test]
fn test_encrypted_pdf_with_empty_passwords() {
    let doc = Document::builder()
        .encryption(PdfEncryption {
            user_password: Some("".into()),
            owner_password: Some("".into()),
            allow_print: true,
            allow_copy: true,
            allow_modify: true,
        })
        .content(DocumentNode::Text(TextNode::new("无密码加密测试")))
        .build();

    let pdf_bytes = PdfMake::new().render(&doc).expect("生成 PDF 失败");
    let text = String::from_utf8_lossy(&pdf_bytes);

    // 空密码不应生成加密字典
    assert!(!text.contains("/Encrypt"));
}

#[test]
fn test_encrypted_pdf_with_only_user_password() {
    let doc = Document::builder()
        .encryption(PdfEncryption {
            user_password: Some("user123".into()),
            owner_password: None,
            allow_print: true,
            allow_copy: false,
            allow_modify: false,
        })
        .content(DocumentNode::Text(TextNode::new("仅用户密码测试")))
        .build();

    let pdf_bytes = PdfMake::new().render(&doc).expect("生成 PDF 失败");
    let text = String::from_utf8_lossy(&pdf_bytes);

    // 应生成加密字典
    assert!(text.contains("/Encrypt"));
}

#[test]
fn test_encrypt_dictionary_uses_hex_strings_for_o_and_u() {
    let doc = Document::builder()
        .encryption(PdfEncryption {
            user_password: Some("123123".into()),
            owner_password: Some("123123".into()),
            allow_print: true,
            allow_copy: false,
            allow_modify: false,
        })
        .content(DocumentNode::Text(TextNode::new("test")))
        .build();

    let pdf_bytes = PdfMake::new().render(&doc).expect("render");
    let text = String::from_utf8_lossy(&pdf_bytes);

    assert!(text.contains("/O <"), "O value should use hex string format");
    assert!(text.contains("/U <"), "U value should use hex string format");
}

#[test]
fn test_encrypt_dictionary_u_value_is_32_bytes() {
    let doc = Document::builder()
        .encryption(PdfEncryption {
            user_password: Some("123123".into()),
            owner_password: Some("123123".into()),
            allow_print: true,
            allow_copy: false,
            allow_modify: false,
        })
        .content(DocumentNode::Text(TextNode::new("test")))
        .build();

    let pdf_bytes = PdfMake::new().render(&doc).expect("render");
    let text = String::from_utf8_lossy(&pdf_bytes);

    // Extract /U hex value: find "/U <" and then the hex chars before ">"
    let u_start = text.find("/U <").expect("should have /U <") + 4;
    let u_end = text[u_start..].find('>').expect("should have closing >") + u_start;
    let u_hex = &text[u_start..u_end];
    assert_eq!(u_hex.len(), 64, "/U hex string should be 64 chars (32 bytes), got {}", u_hex.len());
}

#[test]
fn test_encrypted_pdf_with_only_owner_password() {
    let doc = Document::builder()
        .encryption(PdfEncryption {
            user_password: None,
            owner_password: Some("owner456".into()),
            allow_print: true,
            allow_copy: false,
            allow_modify: false,
        })
        .content(DocumentNode::Text(TextNode::new("仅所有者密码测试")))
        .build();

    let pdf_bytes = PdfMake::new().render(&doc).expect("生成 PDF 失败");
    let text = String::from_utf8_lossy(&pdf_bytes);

    // 应生成加密字典
    assert!(text.contains("/Encrypt"));
}
