use crate::model::{DocumentInfo, PdfEncryption};
use md5::{Md5, Digest};

fn rc4(key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut s: Vec<u8> = (0..=255).collect();
    let mut j: u8 = 0;
    for i in 0..256 {
        j = j.wrapping_add(s[i]).wrapping_add(key[i % key.len()]);
        s.swap(i, j as usize);
    }
    let mut i: u8 = 0;
    let mut j: u8 = 0;
    let mut output = Vec::with_capacity(data.len());
    for &byte in data {
        i = i.wrapping_add(1);
        j = j.wrapping_add(s[i as usize]);
        s.swap(i as usize, j as usize);
        output.push(byte ^ s[(s[i as usize].wrapping_add(s[j as usize])) as usize]);
    }
    output
}

/// 对流数据进行 RC4 加密
///
/// PDF 规范要求对所有流数据进行加密。加密在压缩之后进行。
/// 如果设置了加密密钥，此函数将使用 RC4 算法加密数据。
pub fn rc4_encrypt_stream(key: &[u8], data: &[u8]) -> Vec<u8> {
    if key.is_empty() {
        return data.to_vec();
    }
    rc4(key, data)
}

fn pad_password(pwd: &str) -> Vec<u8> {
    let mut padded = pwd.as_bytes().to_vec();
    padded.truncate(32);
    let padding: &[u8] = &[
        0x28, 0xBF, 0x4E, 0x5E, 0x4E, 0x75, 0x8A, 0x41, 0x64, 0x00, 0x4E, 0x56, 0xFF, 0xFA, 0x01, 0x08,
        0x2E, 0x2E, 0x00, 0xB6, 0xD0, 0x68, 0x3E, 0x80, 0x2F, 0x0C, 0xA9, 0xFE, 0x64, 0x53, 0x69, 0x7A,
    ];
    let remaining = 32 - padded.len();
    padded.extend_from_slice(&padding[..remaining]);
    padded
}

pub fn compute_owner_password(user_pwd: &str, owner_pwd: &str) -> Vec<u8> {
    let padded_owner = pad_password(owner_pwd);
    let mut hasher = Md5::new();
    hasher.update(&padded_owner);
    let key = &hasher.finalize()[..5]; // 40-bit key
    let padded_user = pad_password(user_pwd);
    rc4(key, &padded_user)
}

/// 计算 R=3 的 O 值（owner password value）
///
/// 根据 PDF 规范 v1.4 R=3：
/// 1. 对 owner 密码做 padding
/// 2. MD5 哈希 50 次
/// 3. 取前 5 字节作为 RC4 key
/// 4. 用该 key 加密 padded user password
pub fn compute_owner_password_r3(user_pwd: &str, owner_pwd: &str) -> Vec<u8> {
    let padded_owner = pad_password(owner_pwd);
    let mut hasher = Md5::new();
    hasher.update(&padded_owner);
    let mut digest = hasher.finalize().to_vec();

    for _ in 0..50 {
        let mut hasher = Md5::new();
        hasher.update(&digest[..16]);
        digest = hasher.finalize().to_vec();
    }

    let key = digest[..16].to_vec();
    let padded_user = pad_password(user_pwd);
    let mut result = rc4(&key, &padded_user);
    for i in 1..20u8 {
        let iter_key: Vec<u8> = key.iter().map(|b| b ^ i).collect();
        result = rc4(&iter_key, &result);
    }
    result
}

/// 计算 R=3 的 U 值
///
/// 根据 PDF 规范 v1.4 R=3：
/// 1. MD5(padding + file_id)
/// 2. RC4(hash, encryption_key)
/// 3. 拼接 16 字节任意数据（用 0 填充），总共 32 字节
fn compute_user_password_r3(user_pwd: &str, o_value: &[u8], permissions: i32, file_id: &[u8]) -> Vec<u8> {
    let key = compute_encryption_key_r3(user_pwd, o_value, permissions, file_id);
    let padding: &[u8] = &[
        0x28, 0xBF, 0x4E, 0x5E, 0x4E, 0x75, 0x8A, 0x41, 0x64, 0x00, 0x4E, 0x56, 0xFF, 0xFA, 0x01, 0x08,
        0x2E, 0x2E, 0x00, 0xB6, 0xD0, 0x68, 0x3E, 0x80, 0x2F, 0x0C, 0xA9, 0xFE, 0x64, 0x53, 0x69, 0x7A,
    ];
    let mut hasher = Md5::new();
    hasher.update(padding);
    hasher.update(file_id);
    let hash = hasher.finalize();
    let mut u = rc4(&key, &hash[..16]);
    for i in 1..20u8 {
        let iter_key: Vec<u8> = key.iter().map(|b| b ^ i).collect();
        u = rc4(&iter_key, &u);
    }
    let mut result = Vec::with_capacity(32);
    result.extend_from_slice(&u[..16]);
    result.extend_from_slice(&padding[..16]);
    result
}

/// 计算 R=3 的 128-bit 加密密钥
///
/// 根据 PDF 规范 v1.4 R=3：
/// 1. 对用户密码进行 padding
/// 2. 拼接 O 值、P 值（小端序）、文件 ID
/// 3. 计算 MD5 哈希
/// 4. 再做 50 轮 MD5
/// 5. 取前 16 字节作为 128-bit 密钥
fn compute_encryption_key_r3(
    user_pwd: &str,
    o_value: &[u8],
    permissions: i32,
    file_id: &[u8],
) -> Vec<u8> {
    let padded_user = pad_password(user_pwd);
    let mut data = Vec::new();
    data.extend_from_slice(&padded_user);
    data.extend_from_slice(o_value);
    data.extend_from_slice(&permissions.to_le_bytes());
    data.extend_from_slice(file_id);

    let mut hasher = Md5::new();
    hasher.update(&data);
    let mut hash = hasher.finalize().to_vec();

    for _ in 0..50 {
        let mut hasher = Md5::new();
        hasher.update(&hash[..16]);
        hash = hasher.finalize().to_vec();
    }

    hash[..16].to_vec()
}

/// 计算加密密钥（128-bit）- 公开接口
///
/// 根据 PDF 规范 v1.4 R=3，加密密钥通过以下方式派生：
/// 1. 对用户密码进行 padding
/// 2. 拼接 O 值、P 值（小端序）和文件 ID
/// 3. 计算 MD5 哈希，再做 4 轮额外 MD5
/// 4. 取前 16 字节作为 128-bit 密钥
pub fn compute_encryption_key(
    user_pwd: &str,
    o_value: &[u8],
    permissions: i32,
    file_id: &[u8],
) -> Vec<u8> {
    compute_encryption_key_r3(user_pwd, o_value, permissions, file_id)
}

/// 计算 U 值（用户密码验证值）- 公开接口
///
/// 根据 PDF 规范 v1.4 R=3：
/// U = RC4(MD5(padding + file_id), key)
pub fn compute_user_password_with_id(
    user_pwd: &str,
    o_value: &[u8],
    permissions: i32,
    file_id: &[u8],
) -> Vec<u8> {
    compute_user_password_r3(user_pwd, o_value, permissions, file_id)
}

/// 派生 per-object 加密密钥
///
/// 根据 PDF 规范，每个对象使用独立的加密密钥：
/// key = MD5(global_key + object_number_bytes + generation_bytes)[..10]
fn derive_object_key(global_key: &[u8], object_number: u32, generation: u32) -> Vec<u8> {
    let mut data = Vec::with_capacity(global_key.len() + 5);
    data.extend_from_slice(global_key);
    data.extend_from_slice(&object_number.to_le_bytes()[..3]); // 3 bytes low-endian
    data.push((generation & 0xFF) as u8);
    data.push(((generation >> 8) & 0xFF) as u8);
    let mut hasher = Md5::new();
    hasher.update(&data);
    hasher.finalize()[..16].to_vec()
}

/// 对 PDF 字符串进行加密
///
/// PDF 规范要求加密后，所有字符串对象也必须加密。
/// 此函数将字符串转换为与 pdf_string_value 相同格式的字节，使用 RC4 加密。
pub fn encrypt_pdf_string(key: &[u8], text: &str) -> String {
    if key.is_empty() {
        return pdf_string_value(text);
    }
    // 使用与 pdf_string_value 相同的编码逻辑
    let bytes = if text.is_ascii() {
        text.as_bytes().to_vec()
    } else {
        // 非 ASCII 字符：使用 UTF-16BE 编码（带 BOM）
        let mut utf16_bytes = Vec::new();
        utf16_bytes.extend_from_slice(&[0xFE, 0xFF]); // BOM
        for ch in text.encode_utf16() {
            utf16_bytes.push((ch >> 8) as u8);
            utf16_bytes.push((ch & 0xFF) as u8);
        }
        utf16_bytes
    };
    let encrypted = rc4(key, &bytes);
    escape_pdf_string_binary(&encrypted)
}

/// 使用 per-object key 加密 PDF 字符串
///
/// 先从全局 key 派生对象特定 key，再用 RC4 加密
pub fn encrypt_pdf_string_with_object_key(
    global_key: &[u8],
    object_number: u32,
    text: &str,
) -> String {
    if global_key.is_empty() {
        return pdf_string_value(text);
    }
    let obj_key = derive_object_key(global_key, object_number, 0);
    // 使用与 pdf_string_value 相同的编码逻辑
    let bytes = if text.is_ascii() {
        text.as_bytes().to_vec()
    } else {
        // 非 ASCII 字符：使用 UTF-16BE 编码（带 BOM）
        let mut utf16_bytes = Vec::new();
        utf16_bytes.extend_from_slice(&[0xFE, 0xFF]); // BOM
        for ch in text.encode_utf16() {
            utf16_bytes.push((ch >> 8) as u8);
            utf16_bytes.push((ch & 0xFF) as u8);
        }
        utf16_bytes
    };
    let encrypted = rc4(&obj_key, &bytes);
    escape_pdf_string_binary(&encrypted)
}

fn escape_pdf_string_binary(data: &[u8]) -> String {
    let mut result = Vec::with_capacity(data.len() * 4 + 2);
    result.push(b'(');
    for &byte in data {
        match byte {
            b'(' => result.extend_from_slice(b"\\("),
            b')' => result.extend_from_slice(b"\\)"),
            b'\\' => result.extend_from_slice(b"\\\\"),
            0x0A => result.extend_from_slice(b"\\n"),
            0x0D => result.extend_from_slice(b"\\r"),
            0x09 => result.extend_from_slice(b"\\t"),
            0x08 => result.extend_from_slice(b"\\b"),
            0x0C => result.extend_from_slice(b"\\f"),
            0x20..=0x7E => result.push(byte), // Printable ASCII
            _ => {
                // Use octal escape for non-printable bytes
                result.push(b'\\');
                result.push(b'0' + (byte >> 6));
                result.push(b'0' + ((byte >> 3) & 7));
                result.push(b'0' + (byte & 7));
            }
        }
    }
    result.push(b')');
    String::from_utf8(result).unwrap_or_default()
}

#[derive(Default)]
pub struct PdfWriter {
    objects: Vec<Vec<u8>>,
    encryption_key: Option<Vec<u8>>,
    file_id: Option<Vec<u8>>,
    enc_o_value: Option<Vec<u8>>,
    enc_permissions: i32,
    enc_user_pwd: String,
}

impl PdfWriter {
    pub fn new() -> Self {
        Self::default()
    }

    /// 设置加密上下文（加密密钥和文件 ID）
    pub fn set_encryption_context(&mut self, key: Vec<u8>, file_id: Vec<u8>) {
        self.encryption_key = Some(key);
        self.file_id = Some(file_id);
    }

    /// 获取加密密钥（如果已设置）
    pub fn encryption_key(&self) -> Option<&[u8]> {
        self.encryption_key.as_deref()
    }

    /// 获取文件 ID（如果已设置）
    pub fn file_id(&self) -> Option<&[u8]> {
        self.file_id.as_deref()
    }

    /// 从 PdfEncryption 配置设置完整加密上下文
    ///
    /// 计算 O 值、权限、加密密钥和文件 ID，并存储在 writer 中。
    /// 后续调用 `finish()` 时会使用这些存储的值，避免重复生成。
    pub fn setup_encryption(&mut self, enc: &PdfEncryption) {
        let has_user_pwd = enc.user_password.as_ref().is_some_and(|p| !p.is_empty());
        let has_owner_pwd = enc.owner_password.as_ref().is_some_and(|p| !p.is_empty());

        if has_user_pwd || has_owner_pwd {
            let user_pwd = enc.user_password.as_deref().unwrap_or("");
            let owner_pwd = enc.owner_password.as_deref().unwrap_or("");

            let o_value = compute_owner_password_r3(user_pwd, owner_pwd);

            let mut permissions = -3904i32;
            if !enc.allow_print { permissions &= !4; }
            if !enc.allow_copy { permissions &= !16; }
            if !enc.allow_modify { permissions &= !8; }

            let file_id = self.generate_file_id();
            let encryption_key = compute_encryption_key_r3(user_pwd, &o_value, permissions, &file_id);

            self.encryption_key = Some(encryption_key);
            self.file_id = Some(file_id);
            self.enc_o_value = Some(o_value);
            self.enc_permissions = permissions;
            self.enc_user_pwd = user_pwd.to_string();
        }
    }

    pub fn generate_file_id(&self) -> Vec<u8> {
        use std::time::{SystemTime, UNIX_EPOCH};
        let mut hasher = Md5::new();
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        hasher.update(timestamp.to_le_bytes());
        hasher.update(self.objects.len().to_le_bytes());
        hasher.finalize().to_vec()
    }

    pub fn next_object_id(&self) -> usize {
        self.objects.len() + 1
    }

    pub fn add_object(&mut self, object: Vec<u8>) -> usize {
        let id = self.next_object_id();
        self.objects.push(object);
        id
    }

    pub fn add_object_with_stream(&mut self, dictionary: Vec<u8>, stream: Vec<u8>) -> usize {
        let object_number = self.next_object_id() as u32;
        let mut object = Vec::new();
        object.extend_from_slice(&dictionary);
        object.extend_from_slice(b"\nstream\n");
        
        // 如果设置了加密密钥，使用 per-object key 对流数据进行 RC4 加密
        let final_stream = if let Some(ref key) = self.encryption_key {
            let obj_key = derive_object_key(key, object_number, 0);
            rc4_encrypt_stream(&obj_key, &stream)
        } else {
            stream
        };
        
        object.extend_from_slice(&final_stream);
        object.extend_from_slice(b"\nendstream");
        self.add_object(object)
    }

    pub fn update_object(&mut self, id: usize, data: Vec<u8>) {
        let index = id - 1;
        if index < self.objects.len() {
            self.objects[index] = data;
        }
    }

    pub fn finish(mut self, root_object_id: usize, info: Option<&DocumentInfo>, encryption: Option<&PdfEncryption>, pdf_version: Option<&str>) -> Vec<u8> {
        let version = pdf_version.unwrap_or("1.4");
        let mut pdf = format!("%PDF-{}\n", version).into_bytes();
        let mut offsets = Vec::new();

        // 如果加密上下文尚未设置（render.rs 未调用 setup_encryption），则在此设置
        if self.encryption_key.is_none()
            && let Some(enc) = encryption {
                self.setup_encryption(enc);
            }

        // 使用存储的加密上下文
        let file_id = self.file_id.clone().unwrap_or_else(|| self.generate_file_id());
        let encryption_key = self.encryption_key.clone().unwrap_or_default();
        let enc_o_value = self.enc_o_value.clone().unwrap_or_default();
        let enc_permissions = self.enc_permissions;
        let enc_user_pwd = self.enc_user_pwd.clone();

        for (index, object) in self.objects.iter().enumerate() {
            offsets.push(pdf.len());
            pdf.extend_from_slice(format!("{} 0 obj\n", index + 1).as_bytes());
            pdf.extend_from_slice(object);
            pdf.extend_from_slice(b"\nendobj\n");
        }

        let info_obj_id = self.objects.len() + 1;
        if let Some(info) = info {
            offsets.push(pdf.len());
            let mut info_dict = format!("{} 0 obj\n<<\n", info_obj_id);

            // 使用 info 对象的 per-object key 加密字符串
            let info_key = if !encryption_key.is_empty() {
                derive_object_key(&encryption_key, info_obj_id as u32, 0)
            } else {
                Vec::new()
            };
            let string_key = if info_key.is_empty() { None } else { Some(info_key.as_slice()) };

            if let Some(ref title) = info.title {
                if let Some(key) = string_key {
                    info_dict.push_str(&format!("/Title {}\n", encrypt_pdf_string(key, title)));
                } else {
                    info_dict.push_str(&format!("/Title {}\n", pdf_string_value(title)));
                }
            }
            if let Some(ref author) = info.author {
                if let Some(key) = string_key {
                    info_dict.push_str(&format!("/Author {}\n", encrypt_pdf_string(key, author)));
                } else {
                    info_dict.push_str(&format!("/Author {}\n", pdf_string_value(author)));
                }
            }
            if let Some(ref subject) = info.subject {
                if let Some(key) = string_key {
                    info_dict.push_str(&format!("/Subject {}\n", encrypt_pdf_string(key, subject)));
                } else {
                    info_dict.push_str(&format!("/Subject {}\n", pdf_string_value(subject)));
                }
            }
            if let Some(ref keywords) = info.keywords {
                if let Some(key) = string_key {
                    info_dict.push_str(&format!("/Keywords {}\n", encrypt_pdf_string(key, keywords)));
                } else {
                    info_dict.push_str(&format!("/Keywords {}\n", pdf_string_value(keywords)));
                }
            }
            if let Some(ref creator) = info.creator {
                if let Some(key) = string_key {
                    info_dict.push_str(&format!("/Creator {}\n", encrypt_pdf_string(key, creator)));
                } else {
                    info_dict.push_str(&format!("/Creator {}\n", pdf_string_value(creator)));
                }
            }
            if let Some(ref producer) = info.producer {
                if let Some(key) = string_key {
                    info_dict.push_str(&format!("/Producer {}\n", encrypt_pdf_string(key, producer)));
                } else {
                    info_dict.push_str(&format!("/Producer {}\n", pdf_string_value(producer)));
                }
            }
            info_dict.push_str(">>\nendobj\n");
            pdf.extend_from_slice(info_dict.as_bytes());
        }

        let mut encrypt_obj_id = 0;
        if !encryption_key.is_empty() {
            let u_value = compute_user_password_r3(&enc_user_pwd, &enc_o_value, enc_permissions, &file_id);

            encrypt_obj_id = self.objects.len() + 1 + if info.is_some() { 1 } else { 0 };
            offsets.push(pdf.len());

            let o_hex: String = enc_o_value.iter().map(|b| format!("{:02X}", b)).collect();
            let u_hex: String = u_value.iter().map(|b| format!("{:02X}", b)).collect();
            let encrypt_dict = format!(
                "{} 0 obj\n<< /Filter /Standard /V 2 /R 3 /Length 128 /O <{}> /U <{}> /P {} >>\nendobj\n",
                encrypt_obj_id, o_hex, u_hex, enc_permissions
            );
            pdf.extend_from_slice(encrypt_dict.as_bytes());
        }

        let xref_offset = pdf.len();
        pdf.extend_from_slice(
            format!("xref\n0 {}\n0000000000 65535 f \n", offsets.len() + 1).as_bytes(),
        );
        for offset in &offsets {
            pdf.extend_from_slice(format!("{:010} 00000 n \n", offset).as_bytes());
        }

        let mut trailer = format!(
            "trailer\n<< /Size {} /Root {} 0 R",
            offsets.len() + 1,
            root_object_id
        );
        if info.is_some() {
            trailer.push_str(&format!(" /Info {} 0 R", info_obj_id));
        }
        if encrypt_obj_id > 0 {
            trailer.push_str(&format!(" /Encrypt {} 0 R", encrypt_obj_id));
        }

        // 添加文件 ID — 使用 PDF hex string 格式 (<hex>)
        let id_hex: String = file_id.iter().map(|b| format!("{:02X}", b)).collect();
        trailer.push_str(&format!(" /ID [<{}> <{}>]", id_hex, id_hex));

        trailer.push_str(&format!(" >>\nstartxref\n{}\n%%EOF\n", xref_offset));
        pdf.extend_from_slice(trailer.as_bytes());

        pdf
    }
}

fn pdf_string_value(s: &str) -> String {
    if s.is_ascii() {
        // ASCII: use PDF string literal ()
        format!("({})", s.replace('\\', "\\\\").replace('(', "\\(").replace(')', "\\)"))
    } else {
        // Non-ASCII: use UTF-16BE hex string <FEFF...>
        let mut hex = String::with_capacity(4 + s.len() * 4);
        hex.push_str("FEFF");
        for ch in s.encode_utf16() {
            hex.push_str(&format!("{:04X}", ch));
        }
        format!("<{}>", hex)
    }
}

#[allow(dead_code)]
fn escape_pdf_string(s: &str) -> String {
    // Check if string contains non-ASCII characters
    if s.is_ascii() {
        // Simple ASCII: just escape special chars
        s.replace('\\', "\\\\")
            .replace('(', "\\(")
            .replace(')', "\\)")
    } else {
        // Non-ASCII: use UTF-16BE encoding with BOM
        // Format: <FEFF ...hex encoded UTF-16BE...>
        let mut hex = String::with_capacity(4 + s.len() * 4);
        hex.push_str("FEFF"); // BOM
        for ch in s.encode_utf16() {
            hex.push_str(&format!("{:04X}", ch));
        }
        format!("<{}>", hex)
    }
}
