//! PDF Repair functionality
//!
//! This module provides functionality to repair damaged or malformed PDF files.

use std::collections::HashSet;

use super::error::{QpdfError, Result};

/// Types of issues that can be repaired
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RepairIssue {
    /// Invalid or missing PDF header
    InvalidHeader,
    /// Damaged or missing cross-reference table
    DamagedXref,
    /// Missing or incorrect stream lengths
    StreamLength,
    /// Invalid object numbers
    InvalidObjectNumbers,
    /// Dangling object references
    DanglingReferences,
    /// Missing required objects
    MissingObjects,
    /// Invalid page tree
    InvalidPageTree,
    /// Encoding issues
    EncodingIssues,
    /// Incorrect endstream/endobj markers
    InvalidMarkers,
}

/// Result of a repair operation
#[derive(Debug, Clone)]
pub struct RepairResult {
    /// Issues that were found and repaired
    pub repaired: Vec<RepairIssue>,
    /// Issues that were found but could not be repaired
    pub unrepaired: Vec<(RepairIssue, String)>,
    /// Warnings during repair
    pub warnings: Vec<String>,
    /// Whether the repair was successful overall
    pub success: bool,
}

impl RepairResult {
    /// Create a successful repair result
    pub fn success() -> Self {
        Self {
            repaired: Vec::new(),
            unrepaired: Vec::new(),
            warnings: Vec::new(),
            success: true,
        }
    }

    /// Add a repaired issue
    pub fn add_repaired(&mut self, issue: RepairIssue) {
        self.repaired.push(issue);
    }

    /// Add an unrepaired issue
    pub fn add_unrepaired(&mut self, issue: RepairIssue, reason: String) {
        self.unrepaired.push((issue, reason));
        self.success = false;
    }

    /// Add a warning
    pub fn add_warning(&mut self, warning: String) {
        self.warnings.push(warning);
    }
}

/// Configuration for PDF repair
#[derive(Debug, Clone)]
pub struct RepairConfig {
    /// Attempt to recover cross-reference table
    pub recover_xref: bool,
    /// Fix stream lengths
    pub fix_stream_lengths: bool,
    /// Remove dangling references
    pub remove_dangling_refs: bool,
    /// Rebuild page tree if needed
    pub rebuild_page_tree: bool,
    /// Maximum recovery attempts
    pub max_recovery_attempts: u32,
}

impl Default for RepairConfig {
    fn default() -> Self {
        Self {
            recover_xref: true,
            fix_stream_lengths: true,
            remove_dangling_refs: true,
            rebuild_page_tree: true,
            max_recovery_attempts: 10,
        }
    }
}

/// Repair a PDF file
pub fn repair_pdf(data: &[u8], config: &RepairConfig) -> Result<(Vec<u8>, RepairResult)> {
    let mut result = RepairResult::success();
    let mut output = data.to_vec();

    // Check and fix PDF header
    if !data.starts_with(b"%PDF-") {
        if let Some(repaired) = fix_header(&mut output) {
            result.add_repaired(RepairIssue::InvalidHeader);
        }
    }

    // Try to recover XRef if configured
    if config.recover_xref {
        match recover_xref(&output) {
            Ok(Some(fixed_data)) => {
                output = fixed_data;
                result.add_repaired(RepairIssue::DamagedXref);
            }
            Ok(None) => {
                // XRef is fine, no action needed
            }
            Err(e) => {
                result.add_warning(format!("XRef recovery attempted but failed: {}", e));
            }
        }
    }

    // Fix stream lengths if configured
    if config.fix_stream_lengths {
        match fix_stream_lengths(&output) {
            Ok(Some(fixed_data)) => {
                output = fixed_data;
                result.add_repaired(RepairIssue::StreamLength);
            }
            Ok(None) => {}
            Err(e) => {
                result.add_warning(format!("Stream length fix failed: {}", e));
            }
        }
    }

    // Remove dangling references if configured
    if config.remove_dangling_refs {
        match remove_dangling_references(&output) {
            Ok(Some(fixed_data)) => {
                output = fixed_data;
                result.add_repaired(RepairIssue::DanglingReferences);
            }
            Ok(None) => {}
            Err(e) => {
                result.add_warning(format!("Remove dangling refs failed: {}", e));
            }
        }
    }

    // Rebuild page tree if configured
    if config.rebuild_page_tree {
        match rebuild_page_tree(&output) {
            Ok(Some(fixed_data)) => {
                output = fixed_data;
                result.add_repaired(RepairIssue::InvalidPageTree);
            }
            Ok(None) => {}
            Err(e) => {
                result.add_warning(format!("Rebuild page tree failed: {}", e));
            }
        }
    }

    Ok((output, result))
}

/// Fix missing or incorrect PDF header
fn fix_header(data: &mut Vec<u8>) -> Option<()> {
    // Find where PDF content actually starts
    let pdf_marker = b"%PDF-";
    if let Some(pos) = data.windows(pdf_marker.len()).position(|w| w == pdf_marker) {
        if pos > 0 {
            // Remove garbage before header
            data.drain(0..pos);
            return Some(());
        }
        return None; // Header is fine
    }

    // No header found, prepend one
    let header = b"%PDF-1.7\n";
    let mut new_data = header.to_vec();
    new_data.extend_from_slice(data);
    *data = new_data;
    Some(())
}

/// Try to recover a damaged cross-reference table
fn recover_xref(data: &[u8]) -> Result<Option<Vec<u8>>> {
    // Find all object definitions
    let mut objects = Vec::new();
    let mut pos = 0;

    while pos < data.len() {
        // Look for object definition pattern: "N G obj"
        if let Some(obj_info) = find_object_def(&data[pos..]) {
            objects.push((pos + obj_info.0, obj_info.1, obj_info.2));
            pos += obj_info.0 + 1;
        } else {
            break;
        }
    }

    if objects.is_empty() {
        return Err(QpdfError::Repair("No objects found in PDF".to_string()));
    }

    // Check if XRef needs rebuilding
    if has_valid_xref(data) {
        return Ok(None); // XRef is fine
    }

    // Rebuild XRef table
    let rebuilt_data = rebuild_xref(data, &objects)?;
    Ok(Some(rebuilt_data))
}

/// Find an object definition starting at the given position
fn find_object_def(data: &[u8]) -> Option<(usize, u32, u32)> {
    // Simple pattern matching for "N G obj"
    let obj_marker = b" obj";

    for (i, window) in data.windows(obj_marker.len()).enumerate() {
        if window == obj_marker {
            // Look back for object number and generation
            let prefix = &data[..i];
            let prefix_str = std::str::from_utf8(prefix).ok()?;

            // Find the last two numbers before "obj"
            let parts: Vec<&str> = prefix_str.split_whitespace().rev().take(2).collect();
            if parts.len() == 2 {
                if let (Ok(generation), Ok(obj)) =
                    (parts[0].parse::<u32>(), parts[1].parse::<u32>())
                {
                    return Some((i, obj, generation));
                }
            }
        }
    }
    None
}

/// Check if the PDF has a valid XRef table
fn has_valid_xref(data: &[u8]) -> bool {
    // Look for xref keyword
    data.windows(4).any(|w| w == b"xref") || data.windows(7).any(|w| w == b"/XRef") // XRef stream
}

/// Rebuild the XRef table for the given objects
fn rebuild_xref(data: &[u8], objects: &[(usize, u32, u32)]) -> Result<Vec<u8>> {
    // This is a simplified implementation
    // A full implementation would properly construct the XRef table

    let mut output = data.to_vec();

    // Find startxref location
    if let Some(startxref_pos) = output.windows(9).rposition(|w| w == b"startxref") {
        // Remove old xref and trailer
        output.truncate(startxref_pos);
    }

    // Build new xref table
    let xref_offset = output.len();
    output.extend_from_slice(b"\nxref\n");
    output.extend_from_slice(format!("0 {}\n", objects.len() + 1).as_bytes());
    output.extend_from_slice(b"0000000000 65535 f \n");

    for (offset, _obj, generation) in objects {
        output.extend_from_slice(format!("{:010} {:05} n \n", offset, generation).as_bytes());
    }

    // Add trailer
    output.extend_from_slice(b"trailer\n<<\n");
    output.extend_from_slice(format!("/Size {}\n", objects.len() + 1).as_bytes());
    output.extend_from_slice(b">>\n");
    output.extend_from_slice(b"startxref\n");
    output.extend_from_slice(format!("{}\n", xref_offset).as_bytes());
    output.extend_from_slice(b"%%EOF\n");

    Ok(output)
}

/// Fix incorrect stream lengths by updating /Length values in stream dictionaries
fn fix_stream_lengths(data: &[u8]) -> Result<Option<Vec<u8>>> {
    let stream_marker = b"stream";
    let endstream_marker = b"endstream";

    let mut output = data.to_vec();
    let mut modified = false;

    let mut pos = 0;
    while pos < output.len() {
        if let Some(stream_pos) = output[pos..]
            .windows(stream_marker.len())
            .position(|w| w == stream_marker)
        {
            let abs_stream_pos = pos + stream_pos;

            if let Some(endstream_pos) = output[abs_stream_pos..]
                .windows(endstream_marker.len())
                .position(|w| w == endstream_marker)
            {
                let stream_start = abs_stream_pos + stream_marker.len();
                let content_start = if output.get(stream_start) == Some(&b'\r') {
                    if output.get(stream_start + 1) == Some(&b'\n') {
                        stream_start + 2
                    } else {
                        stream_start + 1
                    }
                } else if output.get(stream_start) == Some(&b'\n') {
                    stream_start + 1
                } else {
                    stream_start
                };

                let actual_length = abs_stream_pos + endstream_pos - content_start;

                if let Some((length_start, length_end)) =
                    find_length_in_dict_before(&output, abs_stream_pos)
                {
                    let old_len_str =
                        std::str::from_utf8(&output[length_start..length_end]).unwrap_or("");
                    let new_len_str = actual_length.to_string();
                    if old_len_str != new_len_str {
                        let old_len = length_end - length_start;
                        if new_len_str.len() <= old_len {
                            let new_bytes = new_len_str.as_bytes();
                            output[length_start..length_start + new_bytes.len()]
                                .copy_from_slice(new_bytes);
                            for i in length_start + new_bytes.len()..length_end {
                                output[i] = b' ';
                            }
                        } else {
                            let mut new_output =
                                Vec::with_capacity(output.len() + new_len_str.len() - old_len);
                            new_output.extend_from_slice(&output[..length_start]);
                            new_output.extend_from_slice(new_len_str.as_bytes());
                            new_output.extend_from_slice(&output[length_end..]);
                            return Ok(fix_stream_lengths(&new_output)?);
                        }
                        modified = true;
                    }
                }

                pos = abs_stream_pos + endstream_pos + endstream_marker.len();
            } else {
                pos += 1;
            }
        } else {
            break;
        }
    }

    if modified { Ok(Some(output)) } else { Ok(None) }
}

/// Find /Length N in the dictionary immediately before stream_pos.
/// Returns (start, end) of the numeric value N.
fn find_length_in_dict_before(data: &[u8], stream_pos: usize) -> Option<(usize, usize)> {
    let length_key = b"/Length";
    let search_start = stream_pos.saturating_sub(2048).max(0);
    let search_range = &data[search_start..stream_pos];

    if let Some(key_pos) = search_range
        .windows(length_key.len())
        .rposition(|w| w == length_key)
    {
        let abs_key_pos = search_start + key_pos;
        let value_start = abs_key_pos + length_key.len();
        let value_start = skip_whitespace_pos(data, value_start);
        if value_start >= data.len() {
            return None;
        }
        let mut value_end = value_start;
        while value_end < data.len()
            && (data[value_end].is_ascii_digit()
                || data[value_end] == b'+'
                || data[value_end] == b'-')
        {
            value_end += 1;
        }
        if value_end > value_start {
            return Some((value_start, value_end));
        }
    }
    None
}

fn skip_whitespace_pos(data: &[u8], mut pos: usize) -> usize {
    while pos < data.len()
        && (data[pos].is_ascii_whitespace() || data[pos] == b'\r' || data[pos] == b'\n')
    {
        pos += 1;
    }
    pos
}

/// Remove dangling object references by replacing undefined N G R with null
pub fn remove_dangling_references(data: &[u8]) -> Result<Option<Vec<u8>>> {
    use std::collections::HashSet;

    let defined_objs = collect_defined_objects(data);
    if defined_objs.is_empty() {
        return Ok(None);
    }

    let refs = find_all_refs(data);
    let mut output = data.to_vec();
    let mut modified = false;

    for (start, end, obj_num, gen_val) in refs.into_iter().rev() {
        if !defined_objs.contains(&(obj_num, gen_val)) {
            output.splice(start..end, b"null".iter().copied());
            modified = true;
        }
    }

    if modified { Ok(Some(output)) } else { Ok(None) }
}

fn find_all_refs(data: &[u8]) -> Vec<(usize, usize, u32, u32)> {
    let mut refs = Vec::new();
    let mut i = 0;
    while i + 4 < data.len() {
        if data[i] == b'R' && (i == 0 || data[i - 1].is_ascii_whitespace()) {
            if let Some((start, num, gen_val)) = parse_ref_before(data, i) {
                refs.push((start, i + 1, num, gen_val));
            }
        }
        i += 1;
    }
    refs
}

fn parse_ref_before(data: &[u8], r_pos: usize) -> Option<(usize, u32, u32)> {
    if r_pos == 0 {
        return None;
    }
    let mut pos = r_pos - 1;
    while pos > 0 && data[pos].is_ascii_whitespace() {
        pos -= 1;
    }
    let gen_end = pos + 1;
    while pos > 0 && data[pos - 1].is_ascii_digit() {
        pos -= 1;
    }
    let gen_str = std::str::from_utf8(&data[pos..gen_end]).ok()?;
    let gen_val: u32 = gen_str.parse().ok()?;
    let mut pos = pos.saturating_sub(1);
    while pos > 0 && data[pos].is_ascii_whitespace() {
        pos -= 1;
    }
    let num_end = pos + 1;
    while pos > 0 && data[pos - 1].is_ascii_digit() {
        pos -= 1;
    }
    let num_str = std::str::from_utf8(&data[pos..num_end]).ok()?;
    let num: u32 = num_str.parse().ok()?;
    Some((pos, num, gen_val))
}

fn collect_defined_objects(data: &[u8]) -> HashSet<(u32, u32)> {
    let mut defined = HashSet::new();
    let obj_marker = b" obj";
    let mut pos = 0;
    while pos < data.len() {
        if let Some(idx) = data[pos..]
            .windows(obj_marker.len())
            .position(|w| w == obj_marker)
        {
            let obj_end = pos + idx;
            if let Some((num, gen_val)) = parse_obj_def_before(&data[..obj_end]) {
                defined.insert((num, gen_val));
            }
            pos = obj_end + obj_marker.len();
        } else {
            break;
        }
    }
    defined
}

fn parse_obj_def_before(data: &[u8]) -> Option<(u32, u32)> {
    let s = std::str::from_utf8(data).ok()?;
    let parts: Vec<&str> = s.split_whitespace().rev().take(2).collect();
    if parts.len() == 2 {
        let gen_val: u32 = parts[0].parse().ok()?;
        let num: u32 = parts[1].parse().ok()?;
        Some((num, gen_val))
    } else {
        None
    }
}

/// Rebuild page tree: find all /Type /Page objects and update /Pages /Kids array
pub fn rebuild_page_tree(data: &[u8]) -> Result<Option<Vec<u8>>> {
    let page_refs = find_page_objects(data);
    if page_refs.is_empty() {
        return Ok(None);
    }

    let (kids_start, kids_end) = find_kids_array_in_pages(data)?;
    let new_kids: String = page_refs
        .iter()
        .map(|(n, g)| format!("{} {} R ", n, g))
        .collect();
    let new_kids_inner = format!("[ {}]", new_kids.trim_end());

    let mut output = data.to_vec();
    let old_len = kids_end - kids_start;
    let new_bytes = new_kids_inner.as_bytes();
    if new_bytes.len() <= old_len {
        output[kids_start..kids_start + new_bytes.len()].copy_from_slice(new_bytes);
        for i in kids_start + new_bytes.len()..kids_end {
            output[i] = b' ';
        }
    } else {
        output.splice(kids_start..kids_end, new_bytes.iter().copied());
    }

    if let Some((count_start, count_end)) = find_count_in_pages(&output, kids_end) {
        let count_str = page_refs.len().to_string();
        let count_bytes = count_str.as_bytes();
        if count_bytes.len() <= count_end - count_start {
            output[count_start..count_start + count_bytes.len()].copy_from_slice(count_bytes);
        }
    }

    Ok(Some(output))
}

fn find_page_objects(data: &[u8]) -> Vec<(u32, u32)> {
    let mut pages = Vec::new();
    let type_marker = b"/Type";
    let page_name = b"/Page";
    let mut pos = 0;
    while pos < data.len() {
        if let Some(idx) = data[pos..]
            .windows(type_marker.len())
            .position(|w| w == type_marker)
        {
            let after_type = pos + idx + type_marker.len();
            let value_start = skip_ws(data, after_type);
            if value_start < data.len()
                && data[value_start] == b'/'
                && data[value_start..].starts_with(page_name)
            {
                if let Some((num, gen_val)) = find_obj_def_containing(data, pos + idx) {
                    pages.push((num, gen_val));
                }
            }
            pos = after_type + 1;
        } else {
            break;
        }
    }
    pages
}

fn skip_ws(data: &[u8], mut pos: usize) -> usize {
    while pos < data.len() && data[pos].is_ascii_whitespace() {
        pos += 1;
    }
    pos
}

fn find_obj_def_containing(data: &[u8], pos: usize) -> Option<(u32, u32)> {
    let search_start = pos.saturating_sub(500);
    let before = &data[search_start..=pos];
    let obj_marker = b" obj";
    before
        .windows(obj_marker.len())
        .rposition(|w| w == obj_marker)
        .and_then(|idx| parse_obj_def_before(&before[..idx]))
        .map(|(n, g)| (n, g))
}

fn find_kids_array_in_pages(data: &[u8]) -> Result<(usize, usize)> {
    let type_marker = b"/Type";
    let pages_name = b"/Pages";
    let kids_key = b"/Kids";
    let mut pos = 0;
    while pos < data.len() {
        if let Some(idx) = data[pos..]
            .windows(type_marker.len())
            .position(|w| w == type_marker)
        {
            let after = pos + idx + type_marker.len();
            let value_start = skip_ws(data, after);
            if value_start < data.len()
                && data[value_start] == b'/'
                && data[value_start..].starts_with(pages_name)
            {
                let search_end = (pos + 2000).min(data.len());
                if let Some(kids_idx) = data[pos..search_end]
                    .windows(kids_key.len())
                    .position(|w| w == kids_key)
                {
                    let key_end = pos + kids_idx + kids_key.len();
                    let value_start = skip_whitespace_pos(data, key_end);
                    if value_start < data.len() && data[value_start] == b'[' {
                        let mut depth = 1;
                        let mut i = value_start + 1;
                        while i < data.len() && depth > 0 {
                            if data[i] == b'[' {
                                depth += 1;
                            } else if data[i] == b']' {
                                depth -= 1;
                            }
                            i += 1;
                        }
                        return Ok((value_start, i));
                    }
                }
            }
            pos = after + 1;
        } else {
            break;
        }
    }
    Err(QpdfError::Repair("No /Pages /Kids array found".to_string()))
}

fn find_count_in_pages(data: &[u8], after_kids: usize) -> Option<(usize, usize)> {
    let count_key = b"/Count";
    let search_end = (after_kids + 200).min(data.len());
    let range = &data[after_kids..search_end];
    range
        .windows(count_key.len())
        .position(|w| w == count_key)
        .map(|idx| {
            let key_end = after_kids + idx + count_key.len();
            let value_start = skip_whitespace_pos(data, key_end);
            let mut value_end = value_start;
            while value_end < data.len() && data[value_end].is_ascii_digit() {
                value_end += 1;
            }
            (value_start, value_end)
        })
}

/// Quick validation of a PDF file
pub fn quick_validate(data: &[u8]) -> bool {
    // Check header
    if !data.starts_with(b"%PDF-") {
        return false;
    }

    // Check for %%EOF at end (allowing for trailing whitespace)
    let eof_marker = b"%%EOF";
    let end_search = &data[data.len().saturating_sub(100)..];
    if !end_search
        .windows(eof_marker.len())
        .any(|w| w == eof_marker)
    {
        return false;
    }

    // Check for at least one object
    if !data.windows(4).any(|w| w == b" obj") {
        return false;
    }

    true
}

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

    #[test]
    fn test_quick_validate_valid() {
        let data = b"%PDF-1.7\n1 0 obj\n<</Type/Catalog>>\nendobj\n%%EOF\n";
        assert!(quick_validate(data));
    }

    #[test]
    fn test_quick_validate_no_header() {
        let data = b"1 0 obj\n<</Type/Catalog>>\nendobj\n%%EOF\n";
        assert!(!quick_validate(data));
    }

    #[test]
    fn test_quick_validate_no_eof() {
        let data = b"%PDF-1.7\n1 0 obj\n<</Type/Catalog>>\nendobj\n";
        assert!(!quick_validate(data));
    }

    #[test]
    fn test_fix_stream_lengths() {
        let data =
            b"%PDF-1.7\n1 0 obj\n<< /Length 999 >>\nstream\nhello\nendstream\nendobj\n%%EOF\n";
        let result = fix_stream_lengths(data).unwrap();
        assert!(result.is_some());
        let fixed = result.unwrap();
        assert!(std::str::from_utf8(&fixed).unwrap().contains("/Length 6"));
    }

    #[test]
    fn test_remove_dangling_references() {
        let data = b"%PDF-1.7\n1 0 obj\n<< /Ref 99 0 R >>\nendobj\n%%EOF\n";
        let result = remove_dangling_references(data).unwrap();
        assert!(result.is_some());
        let fixed = result.unwrap();
        assert!(std::str::from_utf8(&fixed).unwrap().contains("/Ref null"));
    }

    #[test]
    fn test_remove_dangling_references_defined_kept() {
        let data = b"%PDF-1.7\n1 0 obj\n<< /Ref 2 0 R >>\nendobj\n2 0 obj\n42\nendobj\n%%EOF\n";
        let result = remove_dangling_references(data).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn test_repair_pdf_with_config() {
        let data = b"%PDF-1.7\n1 0 obj\n<< /Length 99 >>\nstream\nx\nendstream\nendobj\n%%EOF\n";
        let config = RepairConfig {
            fix_stream_lengths: true,
            ..Default::default()
        };
        let (output, result) = repair_pdf(data, &config).unwrap();
        assert!(result.repaired.contains(&RepairIssue::StreamLength));
        assert!(std::str::from_utf8(&output).unwrap().contains("/Length 2"));
    }

    #[test]
    fn test_rebuild_page_tree() {
        let data = concat!(
            "%PDF-1.7\n",
            "1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n",
            "2 0 obj\n<< /Type /Pages /Kids [ 3 0 R ] /Count 1 >>\nendobj\n",
            "3 0 obj\n<< /Type /Page /Parent 2 0 R >>\nendobj\n",
            "%%EOF\n"
        );
        let result = rebuild_page_tree(data.as_bytes()).unwrap();
        assert!(result.is_some());
        let fixed = result.unwrap();
        let s = std::str::from_utf8(&fixed).unwrap();
        assert!(s.contains("/Kids"));
        assert!(s.contains("3 0 R"));
    }
}
