//! Object Stream Support
//!
//! This module provides functionality to decompress and parse PDF object streams,
//! which are used in PDF 1.5+ to store multiple objects in a compressed stream.

use super::error::{EnhancedError, Result};
use flate2::read::ZlibDecoder;
use std::collections::HashMap;
use std::io::Read;

/// Represents an object stored in an object stream
#[derive(Debug, Clone)]
pub struct ObjectStreamEntry {
    /// Object number
    pub obj_num: i32,
    /// Object data (as string)
    pub data: String,
}

/// Parses objects from an object stream
pub struct ObjectStreamParser {
    /// Mapping from object number to decompressed object data
    objects: HashMap<i32, String>,
}

impl ObjectStreamParser {
    /// Create a new parser and decompress the object stream
    pub fn new(pdf_data: &[u8], obj_stream_num: i32) -> Result<Self> {
        let mut parser = Self {
            objects: HashMap::new(),
        };
        parser.parse_object_stream(pdf_data, obj_stream_num)?;
        Ok(parser)
    }

    /// Parse an object stream and extract all objects
    fn parse_object_stream(&mut self, pdf_data: &[u8], obj_stream_num: i32) -> Result<()> {
        // Find the object stream by its number
        let obj_pattern = format!("{} 0 obj", obj_stream_num);
        let obj_pos = find_pattern(pdf_data, obj_pattern.as_bytes()).ok_or_else(|| {
            EnhancedError::Generic(format!("Object stream {} not found", obj_stream_num))
        })?;

        // Extract the object stream dictionary and data
        let after_obj = &pdf_data[obj_pos..];

        // Find stream keyword
        let stream_pos = find_pattern(after_obj, b"stream")
            .ok_or_else(|| EnhancedError::Generic("Stream keyword not found".into()))?;

        // Parse dictionary to get /N (number of objects) and /First (offset to first object)
        let dict_data = &after_obj[..stream_pos];
        let dict_str = String::from_utf8_lossy(dict_data);

        let n = extract_int_value(&dict_str, "/N")
            .ok_or_else(|| EnhancedError::Generic("/N not found in object stream".into()))?;
        let first = extract_int_value(&dict_str, "/First")
            .ok_or_else(|| EnhancedError::Generic("/First not found in object stream".into()))?;

        // Get stream data
        let stream_start = obj_pos + stream_pos + 6; // "stream" is 6 bytes

        // Skip newline after "stream"
        let data_start = if pdf_data.get(stream_start) == Some(&b'\r')
            && pdf_data.get(stream_start + 1) == Some(&b'\n')
        {
            stream_start + 2
        } else if pdf_data.get(stream_start) == Some(&b'\n')
            || pdf_data.get(stream_start) == Some(&b'\r')
        {
            stream_start + 1
        } else {
            stream_start
        };

        // Find endstream
        let after_stream_start = &pdf_data[data_start..];
        let endstream_pos = find_pattern(after_stream_start, b"endstream")
            .ok_or_else(|| EnhancedError::Generic("endstream not found".into()))?;

        let raw_data = &after_stream_start[..endstream_pos];

        // Check if FlateDecode is used
        let is_flate = dict_str.contains("/FlateDecode") || dict_str.contains("/Fl");

        // Decompress if needed
        let decompressed = if is_flate && !raw_data.is_empty() {
            let mut decoder = ZlibDecoder::new(raw_data);
            let mut output = Vec::new();
            decoder
                .read_to_end(&mut output)
                .map_err(|e| EnhancedError::Generic(format!("Decompression failed: {}", e)))?;
            output
        } else {
            raw_data.to_vec()
        };

        // Parse the object stream content
        // Format: obj_num offset obj_num offset ... (n pairs) followed by object data
        // IMPORTANT: Work with bytes, not strings, because offsets are byte offsets
        // and the data may contain non-UTF8 binary content

        let first_offset = first as usize;
        if first_offset > decompressed.len() {
            return Err(EnhancedError::Generic(format!(
                "/First offset {} exceeds decompressed data length {}",
                first_offset,
                decompressed.len()
            )));
        }

        // Parse the header (n pairs of object_number and offset)
        // The header is typically ASCII, so we can safely convert it
        let header_bytes = &decompressed[..first_offset];
        let header_part = String::from_utf8_lossy(header_bytes);

        // The objects part stays as bytes since offsets are byte-based
        let objects_bytes = &decompressed[first_offset..];

        let mut obj_entries: Vec<(i32, usize)> = Vec::new();
        let mut nums = header_part.split_whitespace();

        for _ in 0..n {
            if let (Some(obj_str), Some(offset_str)) = (nums.next(), nums.next()) {
                if let (Ok(obj_num), Ok(offset)) =
                    (obj_str.parse::<i32>(), offset_str.parse::<usize>())
                {
                    obj_entries.push((obj_num, offset));
                }
            }
        }

        // Extract each object's data using byte offsets
        for i in 0..obj_entries.len() {
            let (obj_num, start_offset) = obj_entries[i];
            let end_offset = if i + 1 < obj_entries.len() {
                obj_entries[i + 1].1
            } else {
                objects_bytes.len()
            };

            if start_offset < objects_bytes.len() {
                let actual_end = end_offset.min(objects_bytes.len());
                let obj_bytes = &objects_bytes[start_offset..actual_end];
                // Convert to string using lossy conversion for any non-UTF8 content
                let obj_data = String::from_utf8_lossy(obj_bytes);
                self.objects.insert(obj_num, obj_data.trim().to_string());
            }
        }

        Ok(())
    }

    /// Get object data by object number
    pub fn get_object(&self, obj_num: i32) -> Option<&String> {
        self.objects.get(&obj_num)
    }

    /// Check if an object exists
    pub fn has_object(&self, obj_num: i32) -> bool {
        self.objects.contains_key(&obj_num)
    }

    /// Get all object numbers
    pub fn object_numbers(&self) -> impl Iterator<Item = &i32> {
        self.objects.keys()
    }
}

/// Find all object streams in a PDF and parse them
pub fn parse_all_object_streams(pdf_data: &[u8]) -> Result<HashMap<i32, String>> {
    let mut all_objects = HashMap::new();

    // Find all /Type /ObjStm patterns to identify object streams
    let mut pos = 0;
    let objstm_pattern = b"/Type/ObjStm";
    let objstm_pattern_space = b"/Type /ObjStm";

    while pos < pdf_data.len() {
        let found_pos = find_pattern(&pdf_data[pos..], objstm_pattern)
            .or_else(|| find_pattern(&pdf_data[pos..], objstm_pattern_space));

        if let Some(rel_pos) = found_pos {
            let abs_pos = pos + rel_pos;

            // Find the object number for this object stream
            if let Some(obj_num) = find_obj_number_before(pdf_data, abs_pos) {
                // Parse this object stream
                match ObjectStreamParser::new(pdf_data, obj_num) {
                    Ok(parser) => {
                        for (&obj_num, data) in &parser.objects {
                            all_objects.insert(obj_num, data.clone());
                        }
                    }
                    Err(e) => {
                        eprintln!(
                            "Warning: Failed to parse object stream {}: {:?}",
                            obj_num, e
                        );
                    }
                }
            }

            pos = abs_pos + 1;
        } else {
            break;
        }
    }

    Ok(all_objects)
}

/// Find a pattern in binary data
fn find_pattern(data: &[u8], pattern: &[u8]) -> Option<usize> {
    if pattern.is_empty() || data.len() < pattern.len() {
        return None;
    }

    for i in 0..=data.len() - pattern.len() {
        if &data[i..i + pattern.len()] == pattern {
            return Some(i);
        }
    }
    None
}

/// Extract integer value from a dictionary string
fn extract_int_value(dict: &str, key: &str) -> Option<i32> {
    if let Some(pos) = dict.find(key) {
        let after_key = &dict[pos + key.len()..];
        let trimmed = after_key.trim_start();

        let mut num_str = String::new();
        for c in trimmed.chars() {
            if c.is_ascii_digit() || (num_str.is_empty() && c == '-') {
                num_str.push(c);
            } else if !num_str.is_empty() {
                break;
            }
        }

        num_str.parse().ok()
    } else {
        None
    }
}

/// Find object number before a position (looking for "N 0 obj" pattern)
fn find_obj_number_before(data: &[u8], pos: usize) -> Option<i32> {
    let search_start = pos.saturating_sub(100);
    let search_region = &data[search_start..pos];

    // Find " obj" or "\nobj" pattern
    let obj_patterns = [b" obj", b"\nobj", b"\robj"];

    let mut last_obj_pos = None;
    for pattern in &obj_patterns {
        for i in 0..search_region.len().saturating_sub(pattern.len()) {
            if &search_region[i..i + pattern.len()] == *pattern {
                last_obj_pos = Some(search_start + i);
            }
        }
    }

    if let Some(obj_pos) = last_obj_pos {
        // Read backwards to find object number
        let before = &data[obj_pos.saturating_sub(20)..obj_pos];

        // Skip backwards past generation number and whitespace to find object number
        let s = String::from_utf8_lossy(before);
        let parts: Vec<&str> = s.split_whitespace().collect();

        // Last two should be "obj_num" "gen_num"
        if parts.len() >= 2 {
            if let Ok(obj_num) = parts[parts.len() - 2].parse::<i32>() {
                return Some(obj_num);
            }
        }
    }

    None
}

/// Extract page count from PDF data by parsing object streams if needed.
/// Properly walks the /Pages tree from the root (referenced by Catalog) when possible.
pub fn get_page_count_from_object_streams(pdf_data: &[u8]) -> Option<i32> {
    // Parse all object streams (may succeed even if we don't find /Type /ObjStm in raw bytes,
    // since we scan for ObjStm objects)
    let all_objects = match parse_all_object_streams(pdf_data) {
        Ok(objs) => objs,
        Err(_) => return None,
    };

    if all_objects.is_empty() {
        return None;
    }

    // Prefer: walk from Catalog -> /Pages root -> /Count
    if let Some(root_pages_obj_num) = find_root_pages_ref(pdf_data, &all_objects) {
        if let Some(pages_data) = all_objects.get(&root_pages_obj_num) {
            if let Some(count) = extract_count_from_object(pages_data) {
                if count > 0 {
                    return Some(count);
                }
            }
        }
    }

    // Fallback: take max /Count from any /Type /Pages object (root has total count)
    let mut max_count: i32 = 0;
    for (_, obj_data) in &all_objects {
        if (obj_data.contains("/Type/Pages") || obj_data.contains("/Type /Pages"))
            && obj_data.contains("/Count")
        {
            if let Some(count) = extract_count_from_object(obj_data) {
                if count > max_count {
                    max_count = count;
                }
            }
        }
    }

    if max_count > 0 {
        Some(max_count)
    } else {
        // Last resort: count /Type /Page leaf objects
        let mut page_count = 0;
        for (_, obj_data) in &all_objects {
            if (obj_data.contains("/Type/Page") || obj_data.contains("/Type /Page"))
                && !obj_data.contains("/Type/Pages")
                && !obj_data.contains("/Type /Pages")
            {
                page_count += 1;
            }
        }

        if page_count > 0 {
            Some(page_count)
        } else {
            None
        }
    }
}

/// Find the root /Pages object number from Catalog (in trailer or ObjStm).
fn find_root_pages_ref(pdf_data: &[u8], all_objects: &HashMap<i32, String>) -> Option<i32> {
    // Try trailer /Root first
    if let Some(root_num) = find_root_from_trailer_or_xref(pdf_data) {
        if let Some(catalog_data) = all_objects.get(&root_num) {
            if let Some(pages_num) = extract_ref_from_object(catalog_data, "/Pages") {
                return Some(pages_num);
            }
        }
    }

    // Find Catalog in ObjStm and get /Pages ref
    for (obj_num, obj_data) in all_objects {
        if obj_data.contains("/Type /Catalog") || obj_data.contains("/Type/Catalog") {
            if let Some(pages_num) = extract_ref_from_object(obj_data, "/Pages") {
                return Some(pages_num);
            }
        }
    }

    None
}

/// Find /Root object number from trailer or xref stream.
fn find_root_from_trailer_or_xref(pdf_data: &[u8]) -> Option<i32> {
    let content = std::str::from_utf8(pdf_data).ok()?;
    // Trailer: /Root N M R
    if let Some(trailer_pos) = content.rfind("trailer") {
        let region = &content[trailer_pos..(trailer_pos + 512).min(content.len())];
        if let Some(root_pos) = region.find("/Root") {
            let after = &region[root_pos + 5..];
            let num: String = after
                .chars()
                .skip_while(|c| c.is_ascii_whitespace())
                .take_while(|c| c.is_ascii_digit())
                .collect();
            if let Ok(n) = num.parse::<i32>() {
                return Some(n);
            }
        }
    }
    // XRef stream dict: /Root N M R
    for pat in ["/Type /XRef", "/Type/XRef"] {
        if let Some(pos) = content.rfind(pat) {
            let start = pos.saturating_sub(2048);
            let end = (pos + 512).min(content.len());
            let region = &content[start..end];
            if let Some(root_pos) = region.find("/Root") {
                let after = &region[root_pos + 5..];
                let num: String = after
                    .chars()
                    .skip_while(|c| c.is_ascii_whitespace())
                    .take_while(|c| c.is_ascii_digit())
                    .collect();
                if let Ok(n) = num.parse::<i32>() {
                    return Some(n);
                }
            }
        }
    }
    None
}

/// Extract object reference number (e.g. /Pages 2 0 R -> 2) from object data.
fn extract_ref_from_object(obj_data: &str, key: &str) -> Option<i32> {
    let pos = obj_data.find(key)?;
    let after = &obj_data[pos + key.len()..];
    let trimmed = after.trim_start();
    let num: String = trimmed.chars().take_while(|c| c.is_ascii_digit()).collect();
    num.parse().ok()
}

/// Extract JPEG images from decompressed ObjStm content.
/// Used when images are inside object streams (PDF v2.0).
/// Works with raw decompressed bytes to preserve binary JPEG data.
pub fn extract_jpeg_images_from_objstm(pdf_data: &[u8]) -> Result<Vec<Vec<u8>>> {
    let mut images = Vec::new();
    let mut pos = 0;
    let objstm_pattern = b"/Type/ObjStm";
    let objstm_pattern_space = b"/Type /ObjStm";

    while pos < pdf_data.len() {
        let found_pos = find_pattern(&pdf_data[pos..], objstm_pattern)
            .or_else(|| find_pattern(&pdf_data[pos..], objstm_pattern_space));

        if let Some(rel_pos) = found_pos {
            let abs_pos = pos + rel_pos;

            if let Some(obj_num) = find_obj_number_before(pdf_data, abs_pos) {
                if let Ok((decompressed, n, first)) = decompress_objstm_raw(pdf_data, obj_num) {
                    for obj_bytes in iter_objstm_object_bytes(&decompressed, n, first) {
                        if let Some(jpeg) = extract_jpeg_from_obj_data(obj_bytes) {
                            images.push(jpeg);
                        }
                    }
                }
            }

            pos = abs_pos + 1;
        } else {
            break;
        }
    }

    Ok(images)
}

/// Decompress an ObjStm and return (raw_bytes, n, first).
fn decompress_objstm_raw(pdf_data: &[u8], obj_stream_num: i32) -> Result<(Vec<u8>, i32, i32)> {
    let obj_pattern = format!("{} 0 obj", obj_stream_num);
    let obj_pos = find_pattern(pdf_data, obj_pattern.as_bytes())
        .ok_or_else(|| EnhancedError::Generic("ObjStm not found".into()))?;

    let after_obj = &pdf_data[obj_pos..];
    let stream_pos = find_pattern(after_obj, b"stream")
        .ok_or_else(|| EnhancedError::Generic("stream not found".into()))?;

    let dict_data = &after_obj[..stream_pos];
    let dict_str = String::from_utf8_lossy(dict_data);

    let n = extract_int_value(&dict_str, "/N")
        .ok_or_else(|| EnhancedError::Generic("/N not found".into()))?;
    let first = extract_int_value(&dict_str, "/First")
        .ok_or_else(|| EnhancedError::Generic("/First not found".into()))?;

    let stream_start = obj_pos + stream_pos + 6;
    let data_start = if pdf_data.get(stream_start) == Some(&b'\r')
        && pdf_data.get(stream_start + 1) == Some(&b'\n')
    {
        stream_start + 2
    } else if pdf_data.get(stream_start) == Some(&b'\n')
        || pdf_data.get(stream_start) == Some(&b'\r')
    {
        stream_start + 1
    } else {
        stream_start
    };

    let after_stream_start = &pdf_data[data_start..];
    let endstream_pos = find_pattern(after_stream_start, b"endstream")
        .ok_or_else(|| EnhancedError::Generic("endstream not found".into()))?;

    let raw_data = &after_stream_start[..endstream_pos];
    let is_flate = dict_str.contains("/FlateDecode") || dict_str.contains("/Fl");

    let decompressed = if is_flate && !raw_data.is_empty() {
        let mut decoder = ZlibDecoder::new(raw_data);
        let mut output = Vec::new();
        decoder
            .read_to_end(&mut output)
            .map_err(|e| EnhancedError::Generic(format!("Decompression failed: {}", e)))?;
        output
    } else {
        raw_data.to_vec()
    };

    Ok((decompressed, n, first))
}

/// Iterate over raw object bytes in decompressed ObjStm.
fn iter_objstm_object_bytes(
    decompressed: &[u8],
    n: i32,
    first: i32,
) -> impl Iterator<Item = &[u8]> {
    let first_offset = first as usize;

    if first_offset >= decompressed.len() {
        return vec![].into_iter();
    }

    let header = &decompressed[..first_offset];
    let objects_bytes = &decompressed[first_offset..];
    let header_str = String::from_utf8_lossy(header);
    let mut nums = header_str.split_whitespace();

    let mut offsets: Vec<usize> = Vec::new();
    for _ in 0..n {
        if let (Some(_obj_str), Some(offset_str)) = (nums.next(), nums.next()) {
            if let Ok(offset) = offset_str.parse::<usize>() {
                offsets.push(offset);
            }
        }
    }
    offsets.push(objects_bytes.len());

    let mut result = Vec::new();
    for i in 0..offsets.len().saturating_sub(1) {
        let start = offsets[i];
        let end = offsets[i + 1].min(objects_bytes.len());
        if start < end {
            result.push(&objects_bytes[start..end]);
        }
    }
    result.into_iter()
}

/// Extract JPEG from raw object bytes.
fn extract_jpeg_from_obj_data(obj_bytes: &[u8]) -> Option<Vec<u8>> {
    let has_image = obj_bytes.windows(14).any(|w| w == b"/Subtype/Image")
        || obj_bytes.windows(15).any(|w| w == b"/Subtype /Image");
    let has_dct = obj_bytes.windows(10).any(|w| w == b"/DCTDecode")
        || obj_bytes.windows(18).any(|w| w == b"/Filter/DCTDecode")
        || obj_bytes.windows(19).any(|w| w == b"/Filter /DCTDecode");
    if !has_image || !has_dct {
        return None;
    }
    let stream_pos = find_pattern(obj_bytes, b"stream")?;
    let mut data_start = stream_pos + 6;
    while data_start < obj_bytes.len()
        && (obj_bytes[data_start] == b'\r' || obj_bytes[data_start] == b'\n')
    {
        data_start += 1;
    }
    let endstream_pos = find_pattern(&obj_bytes[data_start..], b"endstream")?;
    let jpeg_data = &obj_bytes[data_start..data_start + endstream_pos];
    if jpeg_data.len() >= 2 && jpeg_data[0] == 0xFF && jpeg_data[1] == 0xD8 {
        Some(jpeg_data.to_vec())
    } else {
        None
    }
}

/// Extract /Count value from an object data string
fn extract_count_from_object(obj_data: &str) -> Option<i32> {
    // Try /Count N pattern
    if let Some(pos) = obj_data.find("/Count") {
        let after = &obj_data[pos + 6..];
        let trimmed = after.trim_start();

        let mut num_str = String::new();
        for c in trimmed.chars() {
            if c.is_ascii_digit() {
                num_str.push(c);
            } else if !num_str.is_empty() {
                break;
            }
        }

        if !num_str.is_empty() {
            return num_str.parse().ok();
        }
    }
    None
}

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

    #[test]
    fn test_extract_int_value() {
        assert_eq!(extract_int_value("/N 5 /First 20", "/N"), Some(5));
        assert_eq!(extract_int_value("/N 5 /First 20", "/First"), Some(20));
        assert_eq!(extract_int_value("/N5/First20", "/N"), Some(5));
    }

    #[test]
    fn test_find_pattern() {
        let data = b"hello world";
        assert_eq!(find_pattern(data, b"world"), Some(6));
        assert_eq!(find_pattern(data, b"foo"), None);
    }
}
