//! PDF writer inspired by QPDF's QPDFWriter
//!
//! Serializes a complete PDF document to a file, memory buffer, or any `Write` implementor.

use super::error::{QpdfError, Result};
use super::xref_stream::{XRefEntry, XRefEntryType, XRefStreamEncoder};
use crate::pdf::object::{Array, Dict, Name, ObjRef, Object, PdfString};
use aes::cipher::{BlockEncryptMut, KeyIvInit};
use md5::{Digest, Md5};
use std::collections::HashMap;
use std::io::{Seek, Write};

// ============================================================================
// Enums
// ============================================================================

/// Object stream mode
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ObjectStreamMode {
    Disable,
    Preserve,
    Generate,
}

/// Stream data mode
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamDataMode {
    Uncompress,
    Preserve,
    Compress,
}

/// Stream decode level
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamDecodeLevel {
    None,
    Generalized,
    Specialized,
    All,
}

/// Encryption method
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EncryptionMethod {
    None,
    Rc4,
    Aes128,
    Aes256,
}

/// Print permission level (R3+)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PrintPermission {
    None,
    Low,
    Full,
}

/// Modify permission level (R3+)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModifyPermission {
    None,
    Assembly,
    FormFilling,
    AnnotateAndForm,
    All,
}

// ============================================================================
// Encryption parameters
// ============================================================================

/// R2 (PDF 1.3) encryption parameters
#[derive(Debug, Clone)]
pub struct R2EncryptionParams {
    pub user_password: String,
    pub owner_password: String,
    pub allow_print: bool,
    pub allow_modify: bool,
    pub allow_extract: bool,
    pub allow_annotate: bool,
}

/// R3 (PDF 1.4) encryption parameters
#[derive(Debug, Clone)]
pub struct R3EncryptionParams {
    pub user_password: String,
    pub owner_password: String,
    pub allow_accessibility: bool,
    pub allow_extract: bool,
    pub allow_assemble: bool,
    pub allow_annotate_and_form: bool,
    pub allow_form_filling: bool,
    pub allow_modify_other: bool,
    pub allow_print: PrintPermission,
    pub allow_modify: ModifyPermission,
}

/// R4 (PDF 1.5, AES-128) encryption parameters
#[derive(Debug, Clone)]
pub struct R4EncryptionParams {
    pub user_password: String,
    pub owner_password: String,
    pub allow_accessibility: bool,
    pub allow_extract: bool,
    pub allow_assemble: bool,
    pub allow_annotate_and_form: bool,
    pub allow_form_filling: bool,
    pub allow_modify_other: bool,
    pub allow_print: PrintPermission,
    pub allow_modify: ModifyPermission,
}

/// R6 (PDF 1.7, AES-256) encryption parameters
#[derive(Debug, Clone)]
pub struct R6EncryptionParams {
    pub user_password: String,
    pub owner_password: String,
    pub allow_accessibility: bool,
    pub allow_extract: bool,
    pub allow_assemble: bool,
    pub allow_annotate_and_form: bool,
    pub allow_form_filling: bool,
    pub allow_modify_other: bool,
    pub allow_print: PrintPermission,
    pub allow_modify: ModifyPermission,
}

/// Encryption configuration (stored, actual crypto in separate module)
#[derive(Debug, Clone)]
pub enum EncryptionConfig {
    R2(R2EncryptionParams),
    R3(R3EncryptionParams),
    R4(R4EncryptionParams),
    R6(R6EncryptionParams),
}

// ============================================================================
// PdfVersion
// ============================================================================

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PdfVersion {
    pub major: u32,
    pub minor: u32,
    pub extension_level: u32,
}

impl PdfVersion {
    pub fn new(major: u32, minor: u32) -> Self {
        Self {
            major,
            minor,
            extension_level: 0,
        }
    }

    pub fn to_header(&self) -> String {
        format!("%PDF-{}.{}", self.major, self.minor)
    }
}

// ============================================================================
// PdfWriter
// ============================================================================

pub struct PdfWriter {
    objects: Vec<(ObjRef, Object)>,
    root: Option<ObjRef>,
    info: Option<ObjRef>,
    trailer_extra: HashMap<Name, Object>,

    object_stream_mode: ObjectStreamMode,
    stream_data_mode: StreamDataMode,
    decode_level: StreamDecodeLevel,
    compress_streams: bool,
    recompress_flate: bool,
    content_normalization: bool,
    qdf_mode: bool,
    preserve_unreferenced: bool,
    newline_before_endstream: bool,
    deterministic_id: bool,
    static_id: bool,
    suppress_original_ids: bool,
    extra_header_text: Option<String>,

    linearize: bool,

    min_pdf_version: Option<PdfVersion>,
    forced_pdf_version: Option<PdfVersion>,

    encryption: Option<EncryptionConfig>,
    preserve_encryption: bool,

    next_obj_num: i32,
}

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

impl PdfWriter {
    pub fn new() -> Self {
        Self {
            objects: Vec::new(),
            root: None,
            info: None,
            trailer_extra: HashMap::new(),

            object_stream_mode: ObjectStreamMode::Disable,
            stream_data_mode: StreamDataMode::Preserve,
            decode_level: StreamDecodeLevel::None,
            compress_streams: true,
            recompress_flate: false,
            content_normalization: false,
            qdf_mode: false,
            preserve_unreferenced: false,
            newline_before_endstream: true,
            deterministic_id: false,
            static_id: false,
            suppress_original_ids: false,
            extra_header_text: None,

            linearize: false,

            min_pdf_version: None,
            forced_pdf_version: None,

            encryption: None,
            preserve_encryption: false,

            next_obj_num: 1,
        }
    }

    pub fn add_object(&mut self, obj: Object) -> ObjRef {
        let id = ObjRef::new(self.next_obj_num, 0);
        self.next_obj_num += 1;
        self.objects.push((id, obj));
        id
    }

    pub fn add_object_with_id(&mut self, id: ObjRef, obj: Object) {
        if id.num >= self.next_obj_num {
            self.next_obj_num = id.num + 1;
        }
        self.objects.push((id, obj));
    }

    pub fn set_root(&mut self, root: ObjRef) {
        self.root = Some(root);
    }

    pub fn set_info(&mut self, info: ObjRef) {
        self.info = Some(info);
    }

    pub fn allocate_obj_id(&mut self) -> ObjRef {
        let id = ObjRef::new(self.next_obj_num, 0);
        self.next_obj_num += 1;
        id
    }

    pub fn set_object_stream_mode(&mut self, mode: ObjectStreamMode) {
        self.object_stream_mode = mode;
    }

    pub fn set_stream_data_mode(&mut self, mode: StreamDataMode) {
        self.stream_data_mode = mode;
    }

    pub fn set_decode_level(&mut self, level: StreamDecodeLevel) {
        self.decode_level = level;
    }

    pub fn set_compress_streams(&mut self, compress: bool) {
        self.compress_streams = compress;
    }

    pub fn set_recompress_flate(&mut self, recompress: bool) {
        self.recompress_flate = recompress;
    }

    pub fn set_content_normalization(&mut self, normalize: bool) {
        self.content_normalization = normalize;
    }

    pub fn set_qdf_mode(&mut self, qdf: bool) {
        self.qdf_mode = qdf;
    }

    pub fn set_preserve_unreferenced(&mut self, preserve: bool) {
        self.preserve_unreferenced = preserve;
    }

    pub fn set_newline_before_endstream(&mut self, newline: bool) {
        self.newline_before_endstream = newline;
    }

    pub fn set_deterministic_id(&mut self, deterministic: bool) {
        self.deterministic_id = deterministic;
    }

    pub fn set_static_id(&mut self, static_id: bool) {
        self.static_id = static_id;
    }

    pub fn set_suppress_original_ids(&mut self, suppress: bool) {
        self.suppress_original_ids = suppress;
    }

    pub fn set_extra_header_text(&mut self, text: &str) {
        self.extra_header_text = Some(text.to_string());
    }

    pub fn set_linearization(&mut self, linearize: bool) {
        self.linearize = linearize;
    }

    pub fn set_min_pdf_version(&mut self, version: PdfVersion) {
        self.min_pdf_version = Some(version);
    }

    pub fn force_pdf_version(&mut self, version: PdfVersion) {
        self.forced_pdf_version = Some(version);
    }

    pub fn set_preserve_encryption(&mut self, preserve: bool) {
        self.preserve_encryption = preserve;
    }

    pub fn set_r2_encryption(&mut self, params: R2EncryptionParams) {
        self.encryption = Some(EncryptionConfig::R2(params));
    }

    pub fn set_r3_encryption(&mut self, params: R3EncryptionParams) {
        self.encryption = Some(EncryptionConfig::R3(params));
    }

    pub fn set_r4_encryption(&mut self, params: R4EncryptionParams) {
        self.encryption = Some(EncryptionConfig::R4(params));
    }

    pub fn set_r6_encryption(&mut self, params: R6EncryptionParams) {
        self.encryption = Some(EncryptionConfig::R6(params));
    }

    /// Generate the /Encrypt dictionary and add it to the trailer.
    /// Must be called before write() when encryption is configured.
    /// The `id` parameter is the document ID (first element, 16 bytes); use `&[]` for default.
    pub fn apply_encryption(&mut self, id: &[u8]) -> Result<()> {
        let Some(ref enc_config) = self.encryption else {
            return Ok(());
        };

        let id_bytes: Vec<u8> = if id.is_empty() {
            vec![0u8; 16]
        } else {
            id.iter().take(16).copied().collect()
        };

        let enc_dict = match enc_config {
            EncryptionConfig::R2(params) => {
                let o_value = compute_o_value(
                    params.owner_password.as_bytes(),
                    params.user_password.as_bytes(),
                    5,
                );
                let p_value = r2_permission_flags(params);
                let key = compute_encryption_key_r2(
                    params.user_password.as_bytes(),
                    &o_value,
                    p_value,
                    &id_bytes,
                );
                let u_value = compute_u_value(&key, &id_bytes);

                let mut dict = Dict::new();
                dict.insert(Name::new("Filter"), Object::Name(Name::new("Standard")));
                dict.insert(Name::new("V"), Object::Int(1));
                dict.insert(Name::new("R"), Object::Int(2));
                dict.insert(
                    Name::new("O"),
                    Object::String(PdfString::new(hex_string(&o_value))),
                );
                dict.insert(
                    Name::new("U"),
                    Object::String(PdfString::new(hex_string(&u_value))),
                );
                dict.insert(Name::new("P"), Object::Int(p_value as i64));
                dict
            }
            EncryptionConfig::R3(_) | EncryptionConfig::R4(_) | EncryptionConfig::R6(_) => {
                return Err(QpdfError::Structure(
                    "R3/R4/R6 encryption not yet implemented in apply_encryption".to_string(),
                ));
            }
        };

        let enc_ref = self.add_object(Object::Dict(enc_dict));
        self.trailer_extra
            .insert(Name::new("Encrypt"), Object::Ref(enc_ref));

        Ok(())
    }

    pub fn write_to_file(&self, path: &str) -> Result<()> {
        let mut file = std::fs::File::create(path)?;
        self.write(&mut file)
    }

    pub fn write_to_memory(&self) -> Result<Vec<u8>> {
        let mut cursor = std::io::Cursor::new(Vec::new());
        self.write(&mut cursor)?;
        Ok(cursor.into_inner())
    }

    pub fn write<W: Write + Seek>(&self, output: &mut W) -> Result<()> {
        let version = self
            .forced_pdf_version
            .or(self.min_pdf_version)
            .unwrap_or(PdfVersion::new(1, 7));

        // 1. Header
        output.write_all(version.to_header().as_bytes())?;
        output.write_all(b"\n%\xE2\xE3\xCF\xD3\n")?;
        if let Some(ref text) = self.extra_header_text {
            output.write_all(text.as_bytes())?;
        }

        let use_object_streams = self.object_stream_mode == ObjectStreamMode::Generate;
        let (object_offsets, xref_offset, size) = if use_object_streams {
            self.write_objects_with_streams(output)?
        } else {
            self.write_objects_classic(output)?
        };

        let startxref = xref_offset;
        let id = self.generate_id(output, &object_offsets)?;

        output.write_all(b"trailer\n<<\n")?;
        output.write_all(format!("/Size {}\n", size).as_bytes())?;
        if let Some(root) = self.root {
            output.write_all(format!("/Root {} {} R\n", root.num, root.generation).as_bytes())?;
        }
        if let Some(info) = self.info {
            output.write_all(format!("/Info {} {} R\n", info.num, info.generation).as_bytes())?;
        }
        output.write_all(format!("/ID [{} {}]\n", id, id).as_bytes())?;
        for (k, v) in &self.trailer_extra {
            write_object(output, &Object::Name(k.clone()))?;
            output.write_all(b" ")?;
            write_object(output, v)?;
            output.write_all(b"\n")?;
        }
        output.write_all(b">>\n")?;

        output.write_all(b"startxref\n")?;
        output.write_all(format!("{}\n", startxref).as_bytes())?;
        output.write_all(b"%%EOF\n")?;

        Ok(())
    }

    fn write_objects_classic<W: Write + Seek>(
        &self,
        output: &mut W,
    ) -> Result<(Vec<(ObjRef, u64)>, u64, i32)> {
        let mut offsets: Vec<(ObjRef, u64)> = Vec::new();
        let mut xref_entries: Vec<(u32, XRefEntry)> = Vec::new();

        xref_entries.push((0, XRefEntry::free(0, 65535)));

        for (id, obj) in &self.objects {
            let offset = output.stream_position()? as u64;
            offsets.push((*id, offset));
            xref_entries.push((
                id.num as u32,
                XRefEntry::uncompressed(offset, id.generation as u32),
            ));

            output.write_all(format!("{} {} obj\n", id.num, id.generation).as_bytes())?;
            match obj {
                Object::Stream { dict, data } => {
                    let data = self.prepare_stream_data(data)?;
                    let mut dict = dict.clone();
                    dict.insert(Name::new("Length"), Object::Int(data.len() as i64));
                    write_object(output, &Object::Dict(dict))?;
                    output.write_all(b"\nstream\n")?;
                    if self.newline_before_endstream {
                        output.write_all(b"\n")?;
                    }
                    output.write_all(&data)?;
                    output.write_all(b"\nendstream\n")?;
                }
                _ => {
                    write_object(output, obj)?;
                    output.write_all(b"\n")?;
                }
            }
            output.write_all(b"endobj\n")?;
        }

        let xref_offset = output.stream_position()? as u64;
        write_xref_table(output, &xref_entries)?;
        let size = (self.objects.len() + 1) as i32;
        Ok((offsets, xref_offset, size))
    }

    fn write_objects_with_streams<W: Write + Seek>(
        &self,
        output: &mut W,
    ) -> Result<(Vec<(ObjRef, u64)>, u64, i32)> {
        let mut offsets: Vec<(ObjRef, u64)> = Vec::new();
        let mut xref_entries: Vec<(u32, XRefEntry)> = Vec::new();
        xref_entries.push((0, XRefEntry::free(0, 65535)));

        let (stream_objs, regular_objs): (Vec<_>, Vec<_>) = self
            .objects
            .iter()
            .partition(|(_, o)| matches!(o, Object::Stream { .. }));

        for (id, obj) in stream_objs {
            let offset = output.stream_position()? as u64;
            offsets.push((*id, offset));
            xref_entries.push((
                id.num as u32,
                XRefEntry::uncompressed(offset, id.generation as u32),
            ));

            output.write_all(format!("{} {} obj\n", id.num, id.generation).as_bytes())?;
            match obj {
                Object::Stream { dict, data } => {
                    let data = self.prepare_stream_data(data)?;
                    let mut dict = dict.clone();
                    dict.insert(Name::new("Length"), Object::Int(data.len() as i64));
                    write_object(output, &Object::Dict(dict))?;
                    output.write_all(b"\nstream\n")?;
                    if self.newline_before_endstream {
                        output.write_all(b"\n")?;
                    }
                    output.write_all(&data)?;
                    output.write_all(b"\nendstream\n")?;
                }
                _ => {
                    write_object(output, obj)?;
                    output.write_all(b"\n")?;
                }
            }
            output.write_all(b"endobj\n")?;
        }

        let obj_stream_id = self.next_obj_num;
        let xref_stream_id = if regular_objs.is_empty() {
            self.next_obj_num
        } else {
            obj_stream_id + 1
        };
        if !regular_objs.is_empty() {
            let mut index_data = Vec::new();
            let mut object_data = Vec::new();
            for (i, (id, obj)) in regular_objs.iter().enumerate() {
                let start = object_data.len();
                index_data.extend_from_slice(format!("{} {}", id.num, start).as_bytes());
                index_data.push(b' ');
                let mut buf = Vec::new();
                write_object(&mut buf, obj)?;
                object_data.extend_from_slice(&buf);
                object_data.push(b' ');
            }
            let first_len = index_data.len();
            let mut stream_data = index_data;
            stream_data.extend(object_data);

            let compressed = compress_data(&stream_data)?;
            let mut dict = HashMap::new();
            dict.insert(Name::new("Type"), Object::Name(Name::new("ObjStm")));
            dict.insert(Name::new("N"), Object::Int(regular_objs.len() as i64));
            dict.insert(Name::new("First"), Object::Int(first_len as i64));
            dict.insert(Name::new("Length"), Object::Int(compressed.len() as i64));
            dict.insert(Name::new("Filter"), Object::Name(Name::new("FlateDecode")));

            let obj_stream_offset = output.stream_position()? as u64;
            output.write_all(format!("{} 0 obj\n", obj_stream_id).as_bytes())?;
            write_object(output, &Object::Dict(dict))?;
            output.write_all(b"\nstream\n")?;
            output.write_all(&compressed)?;
            output.write_all(b"\nendstream\nendobj\n")?;

            for (i, (id, _)) in regular_objs.iter().enumerate() {
                xref_entries.push((
                    id.num as u32,
                    XRefEntry::compressed(obj_stream_id as u64, i as u32),
                ));
            }
            xref_entries.push((
                obj_stream_id as u32,
                XRefEntry::uncompressed(obj_stream_offset, 0),
            ));
        }

        let xref_offset = output.stream_position()? as u64;
        xref_entries.push((
            xref_stream_id as u32,
            XRefEntry::uncompressed(xref_offset, 0),
        ));
        xref_entries.sort_by_key(|(k, _)| *k);

        let entries: Vec<_> = xref_entries.iter().map(|(a, b)| (*a, b.clone())).collect();
        let (w0, w1, w2) = XRefStreamEncoder::optimal_widths(&entries);
        let encoder = XRefStreamEncoder::new(w0, w1, w2);
        let xref_data = encoder.encode_all(&entries);

        let mut xref_dict = HashMap::new();
        xref_dict.insert(Name::new("Type"), Object::Name(Name::new("XRef")));
        xref_dict.insert(Name::new("Size"), Object::Int((xref_stream_id + 1) as i64));
        xref_dict.insert(
            Name::new("W"),
            Object::Array(vec![
                Object::Int(w0 as i64),
                Object::Int(w1 as i64),
                Object::Int(w2 as i64),
            ]),
        );
        xref_dict.insert(
            Name::new("Index"),
            Object::Array(vec![
                Object::Int(0),
                Object::Int((xref_stream_id + 1) as i64),
            ]),
        );
        output.write_all(format!("{} 0 obj\n", xref_stream_id).as_bytes())?;
        xref_dict.insert(Name::new("Length"), Object::Int(xref_data.len() as i64));
        write_object(output, &Object::Dict(xref_dict))?;
        output.write_all(b"\nstream\n")?;
        output.write_all(&xref_data)?;
        output.write_all(b"\nendstream\nendobj\n")?;

        let size = (xref_stream_id + 1) as i32;
        Ok((offsets, xref_offset, size))
    }

    fn prepare_stream_data(&self, data: &[u8]) -> Result<Vec<u8>> {
        match self.stream_data_mode {
            StreamDataMode::Uncompress => decompress_stream(data),
            StreamDataMode::Preserve => Ok(data.to_vec()),
            StreamDataMode::Compress => {
                if self.compress_streams {
                    compress_data(data)
                } else {
                    Ok(data.to_vec())
                }
            }
        }
    }

    fn generate_id<W: Write>(&self, _output: &mut W, offsets: &[(ObjRef, u64)]) -> Result<String> {
        if self.static_id {
            return Ok("<00000000000000000000000000000000>".to_string());
        }
        if self.deterministic_id {
            let mut hasher = Md5::new();
            hasher.update(b"pdf-write");
            for (id, off) in offsets {
                hasher.update(id.num.to_le_bytes());
                hasher.update(id.generation.to_le_bytes());
                hasher.update(off.to_le_bytes());
            }
            let result = hasher.finalize();
            return Ok(format!("<{:032x}>", result));
        }
        let nanos = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos();
        let pid = std::process::id();
        let mut hasher = Md5::new();
        hasher.update(nanos.to_le_bytes());
        hasher.update(pid.to_le_bytes());
        let result = hasher.finalize();
        Ok(format!("<{:032x}>", result))
    }
}

fn write_xref_table<W: Write>(output: &mut W, entries: &[(u32, XRefEntry)]) -> Result<()> {
    output.write_all(b"xref\n")?;
    output.write_all(format!("0 {}\n", entries.len()).as_bytes())?;
    for (_, entry) in entries {
        let f2 = format!("{:010}", entry.field2);
        let f3 = format!("{:05}", entry.field3);
        let typ = match entry.entry_type {
            XRefEntryType::Free => "f",
            XRefEntryType::Uncompressed | XRefEntryType::Compressed => "n",
        };
        output.write_all(format!("{} {} {} 00000 \n", f2, f3, typ).as_bytes())?;
    }
    Ok(())
}

fn write_object<W: Write>(w: &mut W, obj: &Object) -> Result<usize> {
    let mut buf = Vec::new();
    write_object_impl(&mut buf, obj)?;
    let n = buf.len();
    w.write_all(&buf)?;
    Ok(n)
}

/// Serialize a PDF object to bytes (pub for use by linearization module)
pub(crate) fn serialize_object(obj: &Object) -> Result<Vec<u8>> {
    let mut buf = Vec::new();
    write_object_impl(&mut buf, obj)?;
    Ok(buf)
}

/// Write an indirect object to output (for linearization). Handles streams.
/// Returns the byte offset where the object was written.
pub(crate) fn write_indirect_object<W: Write + Seek>(
    w: &mut W,
    id: ObjRef,
    obj: &Object,
) -> Result<u64> {
    let offset = w.stream_position()?;
    w.write_all(format!("{} {} obj\n", id.num, id.generation).as_bytes())?;
    match obj {
        Object::Stream { dict, data } => {
            let mut dict = dict.clone();
            dict.insert(Name::new("Length"), Object::Int(data.len() as i64));
            write_object(w, &Object::Dict(dict))?;
            w.write_all(b"\nstream\n")?;
            w.write_all(b"\n")?;
            w.write_all(data)?;
            w.write_all(b"\nendstream\n")?;
        }
        _ => {
            write_object(w, obj)?;
            w.write_all(b"\n")?;
        }
    }
    w.write_all(b"endobj\n")?;
    Ok(offset)
}

fn write_object_impl<W: Write>(w: &mut W, obj: &Object) -> Result<()> {
    match obj {
        Object::Null => w.write_all(b"null")?,
        Object::Bool(true) => w.write_all(b"true")?,
        Object::Bool(false) => w.write_all(b"false")?,
        Object::Int(i) => w.write_all(i.to_string().as_bytes())?,
        Object::Real(r) => {
            let s = format_real(*r);
            w.write_all(s.as_bytes())?;
        }
        Object::Name(n) => {
            w.write_all(b"/")?;
            w.write_all(encode_name(n.as_str()).as_bytes())?;
        }
        Object::String(s) => {
            let bytes = s.as_bytes();
            if bytes.iter().all(|b| {
                b.is_ascii() && !matches!(b, 0x28 | 0x29 | 0x5C | 0x0D | 0x0A | 0x09 | 0x08 | 0x0C)
            }) {
                w.write_all(b"(")?;
                for &b in bytes {
                    if b == b'\\' || b == b'(' || b == b')' {
                        w.write_all(&[b'\\', b])?;
                    } else {
                        w.write_all(&[b])?;
                    }
                }
                w.write_all(b")")?;
            } else {
                w.write_all(b"<")?;
                for b in bytes {
                    w.write_all(format!("{:02X}", b).as_bytes())?;
                }
                w.write_all(b">")?;
            }
        }
        Object::Array(arr) => {
            w.write_all(b"[")?;
            for (i, elem) in arr.iter().enumerate() {
                if i > 0 {
                    w.write_all(b" ")?;
                }
                write_object_impl(w, elem)?;
            }
            w.write_all(b"]")?;
        }
        Object::Dict(d) => {
            w.write_all(b"<<")?;
            let mut keys: Vec<_> = d.keys().collect();
            keys.sort_by(|a, b| a.as_str().cmp(b.as_str()));
            for k in keys {
                if let Some(v) = d.get(k) {
                    w.write_all(b" /")?;
                    w.write_all(encode_name(k.as_str()).as_bytes())?;
                    w.write_all(b" ")?;
                    write_object_impl(w, v)?;
                }
            }
            w.write_all(b">>")?;
        }
        Object::Stream { .. } => {
            return Err(QpdfError::Structure(
                "Stream must be written via write(), not write_object".to_string(),
            ));
        }
        Object::Ref(r) => w.write_all(format!("{} {} R", r.num, r.generation).as_bytes())?,
    }
    Ok(())
}

fn format_real(r: f64) -> String {
    if r.fract() == 0.0 && r.abs() < 1e15 {
        format!("{}.0", r as i64)
    } else {
        let s = format!("{:.6}", r);
        s.trim_end_matches('0').trim_end_matches('.').to_string()
    }
}

fn encode_name(s: &str) -> String {
    let mut out = String::new();
    for c in s.chars() {
        match c {
            '#' | ' ' | '/' | '(' | ')' | '<' | '>' | '[' | ']' | '{' | '}' | '%' => {
                out.push_str(&format!("#{:02X}", c as u8));
            }
            c if c.is_ascii() && !c.is_ascii_control() => out.push(c),
            c => out.push_str(&format!("#{:04X}", c as u32)),
        }
    }
    out
}

fn compress_data(data: &[u8]) -> Result<Vec<u8>> {
    use std::io::Write;
    let mut buf = Vec::new();
    let mut encoder = flate2::write::DeflateEncoder::new(&mut buf, flate2::Compression::default());
    encoder.write_all(data)?;
    encoder.finish()?;
    Ok(buf)
}

fn decompress_stream(data: &[u8]) -> Result<Vec<u8>> {
    use std::io::Read;
    let mut decoder = flate2::read::DeflateDecoder::new(data);
    let mut out = Vec::new();
    decoder.read_to_end(&mut out)?;
    Ok(out)
}

// ============================================================================
// PDF Encryption
// ============================================================================

/// Standard PDF password padding bytes (32 bytes).
const PDF_PASSWORD_PADDING: [u8; 32] = [
    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,
];

/// Encode bytes as PDF hex string (e.g. &lt;0123456789ABCDEF&gt;).
fn hex_string(bytes: &[u8]) -> Vec<u8> {
    let mut result = vec![b'<'];
    for &b in bytes {
        result.extend_from_slice(format!("{:02X}", b).as_bytes());
    }
    result.push(b'>');
    result
}

/// R2 permission flags from R2EncryptionParams.
fn r2_permission_flags(params: &R2EncryptionParams) -> i32 {
    let mut p: i32 = 0;
    if params.allow_print {
        p |= 0x04;
    }
    if params.allow_modify {
        p |= 0x08;
    }
    if params.allow_extract {
        p |= 0x10;
    }
    if params.allow_annotate {
        p |= 0x20;
    }
    p
}

/// Pad user/owner password to 32 bytes using PDF standard padding.
pub fn pad_password(password: &[u8]) -> [u8; 32] {
    let mut result = [0u8; 32];
    let copy_len = password.len().min(32);
    result[..copy_len].copy_from_slice(&password[..copy_len]);
    if copy_len < 32 {
        result[copy_len..].copy_from_slice(&PDF_PASSWORD_PADDING[..32 - copy_len]);
    }
    result
}

/// PKCS7 padding for AES block cipher (block size 16).
pub fn pkcs7_pad(data: &[u8], block_size: usize) -> Vec<u8> {
    let pad_len = block_size - (data.len() % block_size);
    let mut result = data.to_vec();
    result.extend(std::iter::repeat(pad_len as u8).take(pad_len));
    result
}

/// RC4 encryption (stream cipher, XOR-based keystream).
///
/// # Security Warning
/// RC4 is cryptographically broken and should only be used for compatibility
/// with legacy PDF encryption (R2/R3). Do not use for new security-sensitive applications.
pub fn rc4_encrypt(data: &[u8], key: &[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 result = Vec::with_capacity(data.len());
    let mut i: u8 = 0;
    j = 0;

    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);
        let k = s[(s[i as usize].wrapping_add(s[j as usize])) as usize];
        result.push(byte ^ k);
    }

    result
}

/// AES-128-CBC encryption. A random 16-byte IV is generated and prepended to the output.
pub fn aes128_cbc_encrypt(data: &[u8], key: &[u8]) -> Result<Vec<u8>> {
    use aes::cipher::{BlockEncryptMut, KeyIvInit};
    use cbc::cipher::block_padding::Pkcs7;

    type Aes128CbcEnc = cbc::Encryptor<aes::Aes128>;

    if key.len() != 16 {
        return Err(QpdfError::Structure(
            "AES-128 requires 16-byte key".to_string(),
        ));
    }

    let mut iv = [0u8; 16];
    getrandom::fill(&mut iv)
        .map_err(|e| QpdfError::Structure(format!("Failed to generate random IV: {:?}", e)))?;
    let cipher = Aes128CbcEnc::new_from_slices(key, &iv)
        .map_err(|e| QpdfError::Structure(format!("AES key/IV error: {:?}", e)))?;

    let block_size = 16;
    let padded_len = ((data.len() + block_size - 1) / block_size) * block_size;
    let mut buf = vec![0u8; padded_len];
    buf[..data.len()].copy_from_slice(data);

    let ciphertext = cipher
        .encrypt_padded_mut::<Pkcs7>(&mut buf, data.len())
        .map_err(|e| QpdfError::Structure(format!("AES encryption error: {:?}", e)))?;

    let mut output = iv.to_vec();
    output.extend_from_slice(ciphertext);
    Ok(output)
}

/// AES-256-CBC encryption. A random 16-byte IV is generated and prepended to the output.
pub fn aes256_cbc_encrypt(data: &[u8], key: &[u8]) -> Result<Vec<u8>> {
    use aes::cipher::{BlockEncryptMut, KeyIvInit};
    use cbc::cipher::block_padding::Pkcs7;

    type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;

    if key.len() != 32 {
        return Err(QpdfError::Structure(
            "AES-256 requires 32-byte key".to_string(),
        ));
    }

    let mut iv = [0u8; 16];
    getrandom::fill(&mut iv)
        .map_err(|e| QpdfError::Structure(format!("Failed to generate random IV: {:?}", e)))?;
    let cipher = Aes256CbcEnc::new_from_slices(key, &iv)
        .map_err(|e| QpdfError::Structure(format!("AES key/IV error: {:?}", e)))?;

    let block_size = 16;
    let padded_len = ((data.len() + block_size - 1) / block_size) * block_size;
    let mut buf = vec![0u8; padded_len];
    buf[..data.len()].copy_from_slice(data);

    let ciphertext = cipher
        .encrypt_padded_mut::<Pkcs7>(&mut buf, data.len())
        .map_err(|e| QpdfError::Structure(format!("AES encryption error: {:?}", e)))?;

    let mut output = iv.to_vec();
    output.extend_from_slice(ciphertext);
    Ok(output)
}

/// Compute encryption key for R2 (PDF 1.3) from user password, O value, P value, and ID.
///
/// # Security Warning
/// R2 encryption uses MD5 and a 40-bit key, both of which are cryptographically weak.
/// This function exists only for compatibility with legacy PDF documents.
pub fn compute_encryption_key_r2(
    password: &[u8],
    o_value: &[u8],
    p_value: i32,
    id: &[u8],
) -> Vec<u8> {
    let padded = pad_password(password);

    let mut hasher = Md5::new();
    hasher.update(&padded);
    hasher.update(o_value);
    hasher.update(p_value.to_le_bytes());
    hasher.update(id);

    let hash = hasher.finalize();
    hash[..5].to_vec()
}

/// Compute O value (encrypted owner password hash) for legacy PDF encryption.
///
/// # Security Warning
/// Uses MD5 + RC4, which are cryptographically weak. For compatibility only.
pub fn compute_o_value(owner_password: &[u8], user_password: &[u8], key_length: usize) -> Vec<u8> {
    let padded_owner = pad_password(owner_password);

    let mut hash = Md5::digest(&padded_owner).to_vec();
    for _ in 0..50 {
        hash = Md5::digest(&hash).to_vec();
    }

    let rc4_key = &hash[..key_length.min(hash.len())];
    let padded_user = pad_password(user_password);
    rc4_encrypt(&padded_user, rc4_key)
}

/// Compute U value (encrypted user password hash) for R2 encryption.
///
/// # Security Warning
/// Uses RC4, which is cryptographically weak. For compatibility only.
pub fn compute_u_value(key: &[u8], _id: &[u8]) -> Vec<u8> {
    let padded = pad_password(&[]);
    rc4_encrypt(&padded, key)
}

/// Compute object encryption key for a specific PDF object.
fn compute_object_encryption_key(
    key: &[u8],
    obj_num: i32,
    gen_num: i32,
    method: &EncryptionMethod,
) -> Vec<u8> {
    let mut hasher = Md5::new();
    hasher.update(key);
    hasher.update(&obj_num.to_le_bytes()[..3]);
    hasher.update(&gen_num.to_le_bytes()[..2]);

    if matches!(method, EncryptionMethod::Aes128 | EncryptionMethod::Aes256) {
        hasher.update(b"sAlT");
    }

    let hash = hasher.finalize();
    let key_len = match method {
        EncryptionMethod::None => 0,
        EncryptionMethod::Rc4 => (key.len() + 5).min(16),
        EncryptionMethod::Aes128 => 16,
        EncryptionMethod::Aes256 => 32,
    };
    hash[..key_len.min(16)].to_vec()
}

/// Encrypt object data based on encryption config.
pub fn encrypt_data(
    data: &[u8],
    key: &[u8],
    obj_num: i32,
    gen_num: i32,
    method: &EncryptionMethod,
) -> Result<Vec<u8>> {
    match method {
        EncryptionMethod::None => Ok(data.to_vec()),
        EncryptionMethod::Rc4 => {
            let obj_key = compute_object_encryption_key(key, obj_num, gen_num, method);
            Ok(rc4_encrypt(data, &obj_key))
        }
        EncryptionMethod::Aes128 => {
            let obj_key = compute_object_encryption_key(key, obj_num, gen_num, method);
            aes128_cbc_encrypt(data, &obj_key)
        }
        EncryptionMethod::Aes256 => {
            let obj_key = compute_object_encryption_key(key, obj_num, gen_num, method);
            aes256_cbc_encrypt(data, &obj_key)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pdf::object::Name;

    #[test]
    fn test_write_minimal_pdf() {
        let mut writer = PdfWriter::new();
        let mut catalog = Dict::new();
        catalog.insert(Name::new("Type"), Object::Name(Name::new("Catalog")));
        let mut pages = Dict::new();
        pages.insert(Name::new("Type"), Object::Name(Name::new("Pages")));
        pages.insert(Name::new("Kids"), Object::Array(vec![]));
        pages.insert(Name::new("Count"), Object::Int(0));

        let pages_ref = writer.add_object(Object::Dict(pages));
        catalog.insert(Name::new("Pages"), Object::Ref(pages_ref));
        let catalog_ref = writer.add_object(Object::Dict(catalog));
        writer.set_root(catalog_ref);

        let pdf = writer.write_to_memory().unwrap();
        assert!(pdf.starts_with(b"%PDF-1.7"));
        assert!(pdf.ends_with(b"%%EOF\n"));
        assert!(pdf.windows(4).any(|w| w == b"xref"));
    }

    #[test]
    fn test_write_objects_and_streams() {
        let mut writer = PdfWriter::new();
        let mut dict = Dict::new();
        dict.insert(Name::new("Length"), Object::Int(5));
        let stream = Object::Stream {
            dict,
            data: b"Hello".to_vec(),
        };
        writer.add_object(stream);
        let pdf = writer.write_to_memory().unwrap();
        assert!(pdf.windows(6).any(|w| w == b"stream"));
        assert!(pdf.windows(9).any(|w| w == b"endstream"));
        assert!(pdf.windows(5).any(|w| w == b"Hello"));
    }

    #[test]
    fn test_object_serialization() {
        let mut buf = Vec::new();
        write_object(&mut buf, &Object::Null).unwrap();
        assert_eq!(buf, b"null");

        buf.clear();
        write_object(&mut buf, &Object::Bool(true)).unwrap();
        assert_eq!(buf, b"true");

        buf.clear();
        write_object(&mut buf, &Object::Int(42)).unwrap();
        assert_eq!(buf, b"42");

        buf.clear();
        write_object(&mut buf, &Object::Name(Name::new("Type"))).unwrap();
        assert_eq!(buf, b"/Type");

        buf.clear();
        write_object(&mut buf, &Object::String(PdfString::new(b"hi".to_vec()))).unwrap();
        assert_eq!(buf, b"(hi)");

        buf.clear();
        write_object(
            &mut buf,
            &Object::Array(vec![Object::Int(1), Object::Int(2)]),
        )
        .unwrap();
        assert_eq!(buf, b"[1 2]");
    }

    #[test]
    fn test_name_encoding() {
        let mut buf = Vec::new();
        write_object(&mut buf, &Object::Name(Name::new("Name With Space"))).unwrap();
        assert!(buf.starts_with(b"/"));
        assert!(buf.windows(3).any(|w| w == b"#20"));
    }

    #[test]
    fn test_multiple_writes_consistency() {
        let mut writer = PdfWriter::new();
        writer.add_object(Object::Int(1));
        let pdf1 = writer.write_to_memory().unwrap();
        let pdf2 = writer.write_to_memory().unwrap();
        assert_eq!(pdf1.len(), pdf2.len());
    }

    // ========================================================================
    // Encryption tests
    // ========================================================================

    #[test]
    fn test_rc4_encrypt_decrypt() {
        let key = b"key";
        let data = b"Hello, World!";
        let encrypted = rc4_encrypt(data, key);
        assert_ne!(encrypted, data);
        let decrypted = rc4_encrypt(&encrypted, key);
        assert_eq!(decrypted, data);
    }

    #[test]
    fn test_rc4_symmetric() {
        let key = b"secret";
        let data = b"test data";
        let enc = rc4_encrypt(data, key);
        let dec = rc4_encrypt(&enc, key);
        assert_eq!(dec, data);
    }

    #[test]
    fn test_aes128_cbc_encrypt() {
        let key = [0u8; 16];
        let data = b"plaintext";
        let encrypted = aes128_cbc_encrypt(data, &key).unwrap();
        assert!(encrypted.len() >= 16 + data.len());
        let enc2 = aes128_cbc_encrypt(data, &key).unwrap();
        assert_ne!(
            encrypted, enc2,
            "Random IVs should produce different ciphertexts"
        );
    }

    #[test]
    fn test_aes256_cbc_encrypt() {
        let key = [0u8; 32];
        let data = b"plaintext";
        let encrypted = aes256_cbc_encrypt(data, &key).unwrap();
        assert!(encrypted.len() >= 16 + data.len());
        let enc2 = aes256_cbc_encrypt(data, &key).unwrap();
        assert_ne!(
            encrypted, enc2,
            "Random IVs should produce different ciphertexts"
        );
    }

    #[test]
    fn test_aes128_wrong_key_length() {
        let key = [0u8; 8];
        let data = b"data";
        assert!(aes128_cbc_encrypt(data, &key).is_err());
    }

    #[test]
    fn test_aes256_wrong_key_length() {
        let key = [0u8; 16];
        let data = b"data";
        assert!(aes256_cbc_encrypt(data, &key).is_err());
    }

    #[test]
    fn test_pad_password() {
        let empty = pad_password(&[]);
        assert_eq!(empty.len(), 32);
        assert_eq!(empty[..], super::PDF_PASSWORD_PADDING[..]);

        let short = pad_password(b"abc");
        assert_eq!(short.len(), 32);
        assert_eq!(&short[..3], b"abc");
        assert_eq!(&short[3..], &super::PDF_PASSWORD_PADDING[..29]);

        let long = pad_password(&[0u8; 40]);
        assert_eq!(long.len(), 32);
        assert_eq!(&long[..], &[0u8; 32]);
    }

    #[test]
    fn test_pkcs7_pad() {
        let data = b"hi";
        let padded = pkcs7_pad(data, 16);
        assert_eq!(padded.len(), 16);
        assert_eq!(&padded[..2], b"hi");
        assert_eq!(padded[2], 14);
        for i in 3..16 {
            assert_eq!(padded[i], 14);
        }

        let exact = pkcs7_pad(&[0u8; 16], 16);
        assert_eq!(exact.len(), 32);
        assert_eq!(exact[16], 16);
    }

    #[test]
    fn test_compute_encryption_key_r2() {
        let password = b"user";
        let o_value = [0u8; 32];
        let p_value = -4i32;
        let id = [0u8; 16];
        let key = compute_encryption_key_r2(password, &o_value, p_value, &id);
        assert_eq!(key.len(), 5);
    }

    #[test]
    fn test_compute_o_value() {
        let owner = b"owner";
        let user = b"user";
        let o = compute_o_value(owner, user, 5);
        assert!(!o.is_empty());
        assert_eq!(o.len(), 32);
    }

    #[test]
    fn test_compute_u_value() {
        let key = [0x42u8; 5];
        let id = [0u8; 16];
        let u = compute_u_value(&key, &id);
        assert_eq!(u.len(), 32);
    }

    #[test]
    fn test_encrypt_data_none() {
        let data = b"data";
        let key = b"key";
        let result = encrypt_data(data, key, 1, 0, &EncryptionMethod::None).unwrap();
        assert_eq!(result, data);
    }

    #[test]
    fn test_encrypt_data_rc4() {
        let data = b"secret content";
        let key = [0u8; 16];
        let encrypted = encrypt_data(data, &key, 5, 0, &EncryptionMethod::Rc4).unwrap();
        assert_ne!(encrypted, data);
        let decrypted = rc4_encrypt(
            &encrypted,
            &compute_object_encryption_key(&key, 5, 0, &EncryptionMethod::Rc4),
        );
        assert_eq!(decrypted, data);
    }

    #[test]
    fn test_apply_encryption_r2() {
        let mut writer = PdfWriter::new();
        writer.set_r2_encryption(R2EncryptionParams {
            user_password: "user".to_string(),
            owner_password: "owner".to_string(),
            allow_print: true,
            allow_modify: false,
            allow_extract: true,
            allow_annotate: false,
        });
        writer.apply_encryption(&[]).unwrap();
        let mut catalog = Dict::new();
        catalog.insert(Name::new("Type"), Object::Name(Name::new("Catalog")));
        let mut pages = Dict::new();
        pages.insert(Name::new("Type"), Object::Name(Name::new("Pages")));
        pages.insert(Name::new("Kids"), Object::Array(vec![]));
        pages.insert(Name::new("Count"), Object::Int(0));
        let pages_ref = writer.add_object(Object::Dict(pages));
        catalog.insert(Name::new("Pages"), Object::Ref(pages_ref));
        let catalog_ref = writer.add_object(Object::Dict(catalog));
        writer.set_root(catalog_ref);

        let pdf = writer.write_to_memory().unwrap();
        assert!(pdf.windows(8).any(|w| w == b"/Encrypt"));
        assert!(pdf.windows(7).any(|w| w == b"/Filter"));
        assert!(pdf.windows(2).any(|w| w == b"/R"));
    }
}
