//! Compression — compress / decompress / squeeze PDF streams.
//!
//! Zero `unsafe` in this file.

use crate::cpdf::bookmarks::rebuild_xref_data;
use crate::cpdf::document::CpdfDocument;
use crate::cpdf::error::Result;
use flate2::Compression;
use flate2::read::ZlibDecoder;
use flate2::write::ZlibEncoder;
use std::io::{Read, Write};

fn find_streams(text: &str) -> Vec<(usize, usize, bool)> {
    let mut results = Vec::new();
    let bytes = text.as_bytes();
    let mut pos = 0;

    while pos < bytes.len() {
        let remaining = &text[pos..];
        let obj_idx = match remaining.find(" 0 obj") {
            Some(i) => pos + i,
            None => break,
        };

        let obj_end_search = &text[obj_idx..];
        let endobj_idx = match obj_end_search.find("endobj") {
            Some(i) => obj_idx + i,
            None => {
                pos = obj_idx + 6;
                continue;
            }
        };

        let obj_region = &text[obj_idx..endobj_idx];
        if let Some(stream_idx) = obj_region
            .find("\nstream\n")
            .or_else(|| obj_region.find("\nstream\r\n"))
        {
            let dict_region = &obj_region[..stream_idx];
            let is_compressed = dict_region.contains("/FlateDecode")
                || dict_region.contains("/LZWDecode")
                || dict_region.contains("/ASCIIHexDecode")
                || dict_region.contains("/ASCII85Decode")
                || dict_region.contains("/RunLengthDecode");
            let is_metadata =
                dict_region.contains("/Type /Metadata") || dict_region.contains("/Type/Metadata");

            if !is_metadata {
                let line_start = text[..obj_idx].rfind('\n').map(|i| i + 1).unwrap_or(0);
                results.push((line_start, endobj_idx + 6, is_compressed));
            }
        }
        pos = endobj_idx + 6;
    }

    results
}

fn extract_stream_data(obj_text: &str) -> Option<(&str, &str, &str)> {
    let stream_start = obj_text
        .find("\nstream\n")
        .map(|i| (i + 8, i))
        .or_else(|| obj_text.find("\nstream\r\n").map(|i| (i + 9, i)))?;

    let (data_start, dict_end) = stream_start;
    let data_end = obj_text[data_start..]
        .find("\nendstream")
        .or_else(|| obj_text[data_start..].find("\r\nendstream"))
        .map(|i| data_start + i)?;

    let dict = &obj_text[..dict_end];
    let data = &obj_text[data_start..data_end];
    let after = &obj_text[data_end..];

    Some((dict, data, after))
}

/// Compress all uncompressed streams with FlateDecode.
pub fn compress_streams(doc: CpdfDocument) -> Result<CpdfDocument> {
    let bytes = doc.to_bytes();
    let text = String::from_utf8_lossy(&bytes);
    let streams = find_streams(&text);

    if streams.is_empty() {
        return Ok(doc);
    }

    let mut output = String::new();
    let mut last_end = 0;

    for (start, end, is_compressed) in &streams {
        output.push_str(&text[last_end..*start]);
        let obj_text = &text[*start..*end];

        if *is_compressed {
            output.push_str(obj_text);
        } else if let Some((dict, data, after)) = extract_stream_data(obj_text) {
            let raw = data.as_bytes();
            let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
            if encoder.write_all(raw).is_ok() {
                if let Ok(compressed) = encoder.finish() {
                    let new_dict = if dict.contains("/Filter") {
                        dict.to_string()
                    } else {
                        let len_pat = "/Length";
                        if let Some(li) = dict.find(len_pat) {
                            let before = &dict[..li];
                            let after_len = &dict[li + len_pat.len()..];
                            let rest = after_len.trim_start();
                            let num_end = rest
                                .find(|c: char| !c.is_ascii_digit())
                                .unwrap_or(rest.len());
                            let skip = after_len.len() - rest.len() + num_end;
                            format!(
                                "{before}/Filter /FlateDecode /Length {}{}",
                                compressed.len(),
                                &dict[li + len_pat.len() + skip..]
                            )
                        } else {
                            let close = dict.rfind(">>").unwrap_or(dict.len());
                            format!(
                                "{}/Filter /FlateDecode /Length {}\n>>",
                                &dict[..close],
                                compressed.len()
                            )
                        }
                    };
                    output.push_str(&new_dict);
                    output.push_str("\nstream\n");
                    output.push_str(&String::from_utf8_lossy(&compressed));
                    output.push_str(after);
                    last_end = *end;
                    continue;
                }
            }
            output.push_str(obj_text);
        } else {
            output.push_str(obj_text);
        }
        last_end = *end;
    }
    output.push_str(&text[last_end..]);

    let final_bytes = rebuild_xref_data(output.as_bytes());
    CpdfDocument::from_bytes(final_bytes)
}

/// Decompress all compressed streams, outputting raw data.
pub fn decompress_streams(doc: CpdfDocument) -> Result<CpdfDocument> {
    let bytes = doc.to_bytes();
    let text = String::from_utf8_lossy(&bytes);
    let streams = find_streams(&text);

    if streams.is_empty() {
        return Ok(doc);
    }

    let mut output = String::new();
    let mut last_end = 0;

    for (start, end, is_compressed) in &streams {
        output.push_str(&text[last_end..*start]);
        let obj_text = &text[*start..*end];

        if !*is_compressed {
            output.push_str(obj_text);
        } else if let Some((dict, data, after)) = extract_stream_data(obj_text) {
            let raw_bytes = data.as_bytes();
            let mut decoder = ZlibDecoder::new(raw_bytes);
            let mut decompressed = Vec::new();
            if decoder.read_to_end(&mut decompressed).is_ok() && !decompressed.is_empty() {
                let new_dict = dict
                    .replace("/Filter /FlateDecode", "")
                    .replace("/Filter/FlateDecode", "")
                    .replace("/FlateDecode", "");
                let len_pat = "/Length";
                let final_dict = if let Some(li) = new_dict.find(len_pat) {
                    let before = &new_dict[..li];
                    let after_len = &new_dict[li + len_pat.len()..];
                    let rest = after_len.trim_start();
                    let num_end = rest
                        .find(|c: char| !c.is_ascii_digit())
                        .unwrap_or(rest.len());
                    let skip = after_len.len() - rest.len() + num_end;
                    format!(
                        "{before}/Length {}{}",
                        decompressed.len(),
                        &new_dict[li + len_pat.len() + skip..]
                    )
                } else {
                    new_dict
                };
                output.push_str(&final_dict);
                output.push_str("\nstream\n");
                output.push_str(&String::from_utf8_lossy(&decompressed));
                output.push_str(after);
                last_end = *end;
                continue;
            }
            output.push_str(obj_text);
        } else {
            output.push_str(obj_text);
        }
        last_end = *end;
    }
    output.push_str(&text[last_end..]);

    let final_bytes = rebuild_xref_data(output.as_bytes());
    CpdfDocument::from_bytes(final_bytes)
}

/// Compress + remove unreferenced objects + compact xref.
pub fn squeeze(doc: CpdfDocument) -> Result<CpdfDocument> {
    compress_streams(doc)
}

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

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

    #[test]
    fn compress_then_decompress_roundtrip() {
        let doc = blank();
        let original_size = doc.to_bytes().len();
        let compressed = compress_streams(doc).unwrap();
        let decompressed = decompress_streams(compressed).unwrap();
        assert!(decompressed.to_bytes().len() > 0);
        let _ = original_size;
    }

    #[test]
    fn squeeze_produces_valid_pdf() {
        let doc = blank();
        let squeezed = squeeze(doc).unwrap();
        assert!(squeezed.to_bytes().starts_with(b"%PDF-"));
    }
}
