pdf_ast/recovery/

diagnostics.rs

use super::*;
use crate::ast::{NodeType, PdfDocument};
use std::collections::HashMap;

/// PDF document diagnostics and health assessment
pub struct DocumentDiagnostics {
    config: DiagnosticsConfig,
    checkers: Vec<Box<dyn HealthChecker>>,
}

/// Configuration for diagnostics
#[derive(Debug, Clone)]
pub struct DiagnosticsConfig {
    pub deep_analysis: bool,
    pub check_integrity: bool,
    pub analyze_structure: bool,
    pub validate_references: bool,
    pub check_streams: bool,
    pub timeout_ms: u64,
}

impl Default for DiagnosticsConfig {
    fn default() -> Self {
        Self {
            deep_analysis: true,
            check_integrity: true,
            analyze_structure: true,
            validate_references: true,
            check_streams: true,
            timeout_ms: 30000, // 30 seconds
        }
    }
}

/// Comprehensive health report for a PDF document
#[derive(Debug, Clone)]
pub struct HealthReport {
    pub overall_health: DocumentHealth,
    pub structure_health: StructureHealth,
    pub integrity_score: f64,
    pub corruption_indicators: Vec<CorruptionIndicator>,
    pub recommendations: Vec<Recommendation>,
    pub detailed_findings: HashMap<String, Finding>,
    pub statistics: DiagnosticStatistics,
}

/// Structure-specific health information
#[derive(Debug, Clone)]
pub struct StructureHealth {
    pub has_valid_header: bool,
    pub has_catalog: bool,
    pub has_pages_tree: bool,
    pub has_valid_xref: bool,
    pub has_trailer: bool,
    pub reference_integrity: f64,
    pub stream_integrity: f64,
}

/// Indicator of potential corruption
#[derive(Debug, Clone)]
pub struct CorruptionIndicator {
    pub indicator_type: CorruptionType,
    pub severity: ErrorSeverity,
    pub location: String,
    pub description: String,
    pub confidence: f64,
}

/// Type of corruption detected
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CorruptionType {
    StructuralDamage,
    DataCorruption,
    MissingComponents,
    InvalidReferences,
    StreamCorruption,
    EncodingIssues,
    IntegrityViolation,
}

/// Recommendation for fixing issues
#[derive(Debug, Clone)]
pub struct Recommendation {
    pub priority: Priority,
    pub action: RecommendedAction,
    pub description: String,
    pub estimated_success_rate: f64,
}

/// Recommended action types
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RecommendedAction {
    StructureRepair,
    ReferenceResolution,
    StreamReconstruction,
    EncodingFix,
    DataRecovery,
    ManualIntervention,
}

/// Priority levels for recommendations
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Priority {
    Low = 1,
    Medium = 2,
    High = 3,
    Critical = 4,
}

/// Detailed finding from diagnostic checks
#[derive(Debug, Clone)]
pub struct Finding {
    pub check_name: String,
    pub status: CheckStatus,
    pub details: String,
    pub metrics: HashMap<String, f64>,
}

/// Status of a diagnostic check
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CheckStatus {
    Passed,
    Warning,
    Failed,
    Error,
    Skipped,
}

/// Statistics from diagnostic analysis
#[derive(Debug, Clone, Default)]
pub struct DiagnosticStatistics {
    pub checks_performed: usize,
    pub checks_passed: usize,
    pub checks_failed: usize,
    pub warnings_generated: usize,
    pub analysis_time_ms: u64,
    pub nodes_analyzed: usize,
    pub bytes_analyzed: u64,
}

impl DocumentDiagnostics {
    /// Create a new diagnostics analyzer
    pub fn new(config: DiagnosticsConfig) -> Self {
        let mut diagnostics = Self {
            config: config.clone(),
            checkers: Vec::new(),
        };

        diagnostics.initialize_checkers(&config);
        diagnostics
    }

    /// Perform comprehensive health analysis
    pub fn analyze_health(&self, document: &PdfDocument, data: &[u8]) -> HealthReport {
        let start_time = std::time::Instant::now();
        let mut report = HealthReport {
            overall_health: DocumentHealth::Healthy,
            structure_health: StructureHealth::default(),
            integrity_score: 1.0,
            corruption_indicators: Vec::new(),
            recommendations: Vec::new(),
            detailed_findings: HashMap::new(),
            statistics: DiagnosticStatistics::default(),
        };

        // Analyze structure
        if self.config.analyze_structure {
            self.analyze_structure(document, data, &mut report);
        }

        // Check integrity
        if self.config.check_integrity {
            self.check_integrity(document, data, &mut report);
        }

        // Validate references
        if self.config.validate_references {
            self.validate_references(document, &mut report);
        }

        // Check streams
        if self.config.check_streams {
            self.check_streams(document, data, &mut report);
        }

        // Run all health checkers
        for checker in &self.checkers {
            let finding = checker.check_health(document, data);
            report
                .detailed_findings
                .insert(checker.name().to_string(), finding);
        }

        // Calculate overall health
        self.calculate_overall_health(&mut report);

        // Generate recommendations
        self.generate_recommendations(&mut report);

        // Finalize statistics
        let elapsed = start_time.elapsed().as_millis() as u64;
        report.statistics.analysis_time_ms = elapsed;
        report.statistics.nodes_analyzed = document.ast.get_all_nodes().len();
        report.statistics.bytes_analyzed = data.len() as u64;

        report
    }

    /// Initialize health checkers based on configuration
    fn initialize_checkers(&mut self, config: &DiagnosticsConfig) {
        self.checkers.push(Box::new(HeaderChecker::new()));
        self.checkers.push(Box::new(StructureChecker::new()));
        self.checkers.push(Box::new(ReferenceChecker::new()));

        if config.check_streams {
            self.checkers.push(Box::new(StreamChecker::new()));
        }

        if config.check_integrity {
            self.checkers.push(Box::new(IntegrityChecker::new()));
        }
    }

    /// Analyze document structure
    fn analyze_structure(&self, document: &PdfDocument, data: &[u8], report: &mut HealthReport) {
        let structure = StructureHealth {
            has_valid_header: data.starts_with(b"%PDF-"),
            has_catalog: document.ast.get_root().is_some(),
            has_pages_tree: !document.ast.get_nodes_by_type(NodeType::Pages).is_empty(),
            has_valid_xref: data.windows(4).any(|w| w == b"xref"),
            has_trailer: data.windows(7).any(|w| w == b"trailer"),
            reference_integrity: self.calculate_reference_integrity(document),
            stream_integrity: self.calculate_stream_integrity(document, data),
        };

        // Check conditions before moving structure
        let passed_checks =
            if structure.has_valid_header && structure.has_catalog && structure.has_pages_tree {
                3
            } else {
                0
            };

        report.structure_health = structure;
        report.statistics.checks_performed += 5; // 5 structure checks
        report.statistics.checks_passed += passed_checks;
    }

    /// Check document integrity
    fn check_integrity(&self, document: &PdfDocument, data: &[u8], report: &mut HealthReport) {
        let mut integrity_issues = 0;
        let mut total_checks = 0;

        // Check for truncated data
        total_checks += 1;
        if !data.ends_with(b"%%EOF") {
            integrity_issues += 1;
            report.corruption_indicators.push(CorruptionIndicator {
                indicator_type: CorruptionType::StructuralDamage,
                severity: ErrorSeverity::Warning,
                location: "End of file".to_string(),
                description: "Missing or corrupted EOF marker".to_string(),
                confidence: 0.9,
            });
        }

        // Check for null bytes in inappropriate places
        total_checks += 1;
        if self.contains_inappropriate_nulls(data) {
            integrity_issues += 1;
            report.corruption_indicators.push(CorruptionIndicator {
                indicator_type: CorruptionType::DataCorruption,
                severity: ErrorSeverity::Warning,
                location: "Throughout document".to_string(),
                description: "Null bytes found in inappropriate locations".to_string(),
                confidence: 0.7,
            });
        }

        // Check object count consistency
        total_checks += 1;
        let expected_objects = self.count_object_declarations(data);
        let actual_objects = document.ast.get_all_nodes().len();
        if expected_objects > 0 && actual_objects < expected_objects / 2 {
            integrity_issues += 1;
            report.corruption_indicators.push(CorruptionIndicator {
                indicator_type: CorruptionType::MissingComponents,
                severity: ErrorSeverity::Error,
                location: "Object count".to_string(),
                description: format!(
                    "Expected {} objects, found {}",
                    expected_objects, actual_objects
                ),
                confidence: 0.8,
            });
        }

        // Calculate integrity score
        report.integrity_score = if total_checks > 0 {
            1.0 - (integrity_issues as f64 / total_checks as f64)
        } else {
            1.0
        };

        report.statistics.checks_performed += total_checks;
        report.statistics.checks_passed += total_checks - integrity_issues;
        report.statistics.checks_failed += integrity_issues;
    }

    /// Validate object references
    fn validate_references(&self, document: &PdfDocument, report: &mut HealthReport) {
        let nodes = document.ast.get_all_nodes();
        let mut total_refs = 0;
        let mut broken_refs = 0;

        for node in &nodes {
            let refs = self.extract_references(&node.value);
            total_refs += refs.len();

            for reference in refs {
                if !self.reference_exists(document, &reference) {
                    broken_refs += 1;
                }
            }
        }

        if broken_refs > 0 {
            report.corruption_indicators.push(CorruptionIndicator {
                indicator_type: CorruptionType::InvalidReferences,
                severity: if broken_refs > total_refs / 2 {
                    ErrorSeverity::Critical
                } else {
                    ErrorSeverity::Warning
                },
                location: "Object references".to_string(),
                description: format!("{} broken references out of {}", broken_refs, total_refs),
                confidence: 0.95,
            });
        }

        report.statistics.checks_performed += 1;
        if broken_refs == 0 {
            report.statistics.checks_passed += 1;
        } else {
            report.statistics.checks_failed += 1;
        }
    }

    /// Check stream integrity
    fn check_streams(&self, _document: &PdfDocument, data: &[u8], report: &mut HealthReport) {
        let streams = self.find_streams_in_data(data);
        let mut corrupted_streams = 0;

        for stream in streams {
            if self.is_stream_corrupted(&stream) {
                corrupted_streams += 1;
            }
        }

        if corrupted_streams > 0 {
            report.corruption_indicators.push(CorruptionIndicator {
                indicator_type: CorruptionType::StreamCorruption,
                severity: ErrorSeverity::Warning,
                location: "Stream objects".to_string(),
                description: format!("{} corrupted streams detected", corrupted_streams),
                confidence: 0.8,
            });
        }

        report.statistics.checks_performed += 1;
        if corrupted_streams == 0 {
            report.statistics.checks_passed += 1;
        } else {
            report.statistics.checks_failed += 1;
        }
    }

    /// Calculate overall document health
    fn calculate_overall_health(&self, report: &mut HealthReport) {
        let mut health_score = report.integrity_score;
        let structure = &report.structure_health;

        // Adjust score based on structure
        if !structure.has_valid_header {
            health_score -= 0.2;
        }
        if !structure.has_catalog {
            health_score -= 0.3;
        }
        if !structure.has_pages_tree {
            health_score -= 0.2;
        }
        if !structure.has_valid_xref {
            health_score -= 0.1;
        }
        if !structure.has_trailer {
            health_score -= 0.1;
        }

        // Adjust based on corruption indicators
        let critical_count = report
            .corruption_indicators
            .iter()
            .filter(|i| i.severity == ErrorSeverity::Critical)
            .count();
        let error_count = report
            .corruption_indicators
            .iter()
            .filter(|i| i.severity == ErrorSeverity::Error)
            .count();

        health_score -= critical_count as f64 * 0.2;
        health_score -= error_count as f64 * 0.1;

        // Determine overall health
        report.overall_health = if health_score >= 0.9 {
            DocumentHealth::Healthy
        } else if health_score >= 0.7 {
            DocumentHealth::PartiallyRecovered
        } else if health_score >= 0.4 {
            DocumentHealth::Damaged
        } else {
            DocumentHealth::SeverelyDamaged
        };
    }

    /// Generate recommendations based on findings
    fn generate_recommendations(&self, report: &mut HealthReport) {
        for indicator in &report.corruption_indicators {
            let recommendation = match indicator.indicator_type {
                CorruptionType::StructuralDamage => Recommendation {
                    priority: Priority::High,
                    action: RecommendedAction::StructureRepair,
                    description: "Repair basic PDF structure".to_string(),
                    estimated_success_rate: 0.8,
                },
                CorruptionType::InvalidReferences => Recommendation {
                    priority: Priority::Medium,
                    action: RecommendedAction::ReferenceResolution,
                    description: "Fix broken object references".to_string(),
                    estimated_success_rate: 0.7,
                },
                CorruptionType::StreamCorruption => Recommendation {
                    priority: Priority::Medium,
                    action: RecommendedAction::StreamReconstruction,
                    description: "Reconstruct corrupted streams".to_string(),
                    estimated_success_rate: 0.6,
                },
                CorruptionType::EncodingIssues => Recommendation {
                    priority: Priority::Low,
                    action: RecommendedAction::EncodingFix,
                    description: "Fix text encoding issues".to_string(),
                    estimated_success_rate: 0.9,
                },
                _ => Recommendation {
                    priority: Priority::Medium,
                    action: RecommendedAction::DataRecovery,
                    description: "Attempt general data recovery".to_string(),
                    estimated_success_rate: 0.5,
                },
            };

            report.recommendations.push(recommendation);
        }

        // Sort recommendations by priority
        report
            .recommendations
            .sort_by(|a, b| b.priority.cmp(&a.priority));
    }

    // Helper methods
    fn calculate_reference_integrity(&self, document: &PdfDocument) -> f64 {
        let nodes = document.ast.get_all_nodes();
        if nodes.is_empty() {
            return 1.0;
        }

        let mut total_refs = 0;
        let mut valid_refs = 0;

        for node in &nodes {
            let refs = self.extract_references(&node.value);
            total_refs += refs.len();

            for reference in refs {
                if self.reference_exists(document, &reference) {
                    valid_refs += 1;
                }
            }
        }

        if total_refs == 0 {
            1.0
        } else {
            valid_refs as f64 / total_refs as f64
        }
    }

    fn calculate_stream_integrity(&self, _document: &PdfDocument, data: &[u8]) -> f64 {
        let streams = self.find_streams_in_data(data);
        if streams.is_empty() {
            return 1.0;
        }

        let mut valid_streams = 0;
        for stream in &streams {
            if !self.is_stream_corrupted(stream) {
                valid_streams += 1;
            }
        }

        valid_streams as f64 / streams.len() as f64
    }

    fn contains_inappropriate_nulls(&self, data: &[u8]) -> bool {
        // Check for null bytes in text content (simplified)
        let text_regions = self.find_text_regions(data);
        for region in text_regions {
            if region.contains(&0u8) {
                return true;
            }
        }
        false
    }

    fn count_object_declarations(&self, data: &[u8]) -> usize {
        let data_str = String::from_utf8_lossy(data);
        data_str.matches(" obj").count()
    }

    #[allow(clippy::only_used_in_recursion)]
    fn extract_references(&self, value: &crate::types::PdfValue) -> Vec<String> {
        let mut refs = Vec::new();

        match value {
            crate::types::PdfValue::Reference(r) => {
                refs.push(format!(
                    "{} {} R",
                    r.object_id().number,
                    r.object_id().generation
                ));
            }
            crate::types::PdfValue::Dictionary(dict) => {
                for (_, v) in dict.iter() {
                    refs.extend(self.extract_references(v));
                }
            }
            crate::types::PdfValue::Array(arr) => {
                for v in arr.iter() {
                    refs.extend(self.extract_references(v));
                }
            }
            _ => {}
        }

        refs
    }

    fn reference_exists(&self, document: &PdfDocument, _reference: &str) -> bool {
        // Simplified reference checking
        // In practice, would parse the reference and check if object exists
        !document.ast.get_all_nodes().is_empty()
    }

    fn find_streams_in_data(&self, data: &[u8]) -> Vec<StreamInfo> {
        let mut streams = Vec::new();
        let mut pos = 0;

        while let Some(start) = self.find_pattern(&data[pos..], b"stream") {
            let abs_start = pos + start;
            if let Some(end) = self.find_pattern(&data[abs_start..], b"endstream") {
                let abs_end = abs_start + end;
                streams.push(StreamInfo {
                    start: abs_start,
                    end: abs_end,
                    data: data[abs_start..abs_end].to_vec(),
                });
                pos = abs_end;
            } else {
                pos = abs_start + 6;
            }
        }

        streams
    }

    fn is_stream_corrupted(&self, stream: &StreamInfo) -> bool {
        // Check for common stream corruption indicators
        let data = &stream.data;

        // Check if stream starts properly
        if !data.starts_with(b"stream") {
            return true;
        }

        // Check if stream ends properly
        if !data.ends_with(b"endstream") {
            return true;
        }

        // Check for unexpected null bytes
        let content_start = 6; // Skip "stream"
        let content_end = data.len() - 9; // Skip "endstream"
        if content_end > content_start {
            let content = &data[content_start..content_end];
            // Allow some null bytes but not excessive amounts
            let null_count = content.iter().filter(|&&b| b == 0).count();
            if null_count > content.len() / 4 {
                return true;
            }
        }

        false
    }

    fn find_text_regions<'a>(&self, data: &'a [u8]) -> Vec<&'a [u8]> {
        // Simplified text region detection
        // In practice, would analyze the PDF structure to find text content
        vec![data] // Return entire data for simplification
    }

    fn find_pattern(&self, data: &[u8], pattern: &[u8]) -> Option<usize> {
        data.windows(pattern.len())
            .position(|window| window == pattern)
    }
}

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

impl Default for StructureHealth {
    fn default() -> Self {
        Self {
            has_valid_header: false,
            has_catalog: false,
            has_pages_tree: false,
            has_valid_xref: false,
            has_trailer: false,
            reference_integrity: 0.0,
            stream_integrity: 0.0,
        }
    }
}

#[derive(Debug, Clone)]
#[allow(dead_code)]
struct StreamInfo {
    start: usize,
    end: usize,
    data: Vec<u8>,
}

/// Base trait for health checkers
pub trait HealthChecker: Send + Sync {
    fn name(&self) -> &str;
    fn check_health(&self, document: &PdfDocument, data: &[u8]) -> Finding;
}

/// Header health checker
pub struct HeaderChecker;

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

impl HeaderChecker {
    pub fn new() -> Self {
        Self
    }
}

impl HealthChecker for HeaderChecker {
    fn name(&self) -> &str {
        "HeaderChecker"
    }

    fn check_health(&self, _document: &PdfDocument, data: &[u8]) -> Finding {
        let has_header = data.starts_with(b"%PDF-");
        let mut metrics = HashMap::new();
        metrics.insert("has_header".to_string(), if has_header { 1.0 } else { 0.0 });

        Finding {
            check_name: "Header Validation".to_string(),
            status: if has_header {
                CheckStatus::Passed
            } else {
                CheckStatus::Failed
            },
            details: if has_header {
                "Valid PDF header found".to_string()
            } else {
                "Missing or invalid PDF header".to_string()
            },
            metrics,
        }
    }
}

/// Structure health checker
pub struct StructureChecker;

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

impl StructureChecker {
    pub fn new() -> Self {
        Self
    }
}

impl HealthChecker for StructureChecker {
    fn name(&self) -> &str {
        "StructureChecker"
    }

    fn check_health(&self, document: &PdfDocument, _data: &[u8]) -> Finding {
        let has_root = document.ast.get_root().is_some();
        let node_count = document.ast.get_all_nodes().len();

        let mut metrics = HashMap::new();
        metrics.insert("has_root".to_string(), if has_root { 1.0 } else { 0.0 });
        metrics.insert("node_count".to_string(), node_count as f64);

        let status = if has_root && node_count > 0 {
            CheckStatus::Passed
        } else if has_root {
            CheckStatus::Warning
        } else {
            CheckStatus::Failed
        };

        Finding {
            check_name: "Structure Validation".to_string(),
            status,
            details: format!("Document has {} nodes, root: {}", node_count, has_root),
            metrics,
        }
    }
}

/// Reference health checker
pub struct ReferenceChecker;

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

impl ReferenceChecker {
    pub fn new() -> Self {
        Self
    }
}

impl HealthChecker for ReferenceChecker {
    fn name(&self) -> &str {
        "ReferenceChecker"
    }

    fn check_health(&self, document: &PdfDocument, _data: &[u8]) -> Finding {
        let nodes = document.ast.get_all_nodes();
        let mut metrics = HashMap::new();
        metrics.insert("total_nodes".to_string(), nodes.len() as f64);

        Finding {
            check_name: "Reference Validation".to_string(),
            status: CheckStatus::Passed, // Simplified
            details: "Reference integrity check completed".to_string(),
            metrics,
        }
    }
}

/// Stream health checker
pub struct StreamChecker;

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

impl StreamChecker {
    pub fn new() -> Self {
        Self
    }
}

impl HealthChecker for StreamChecker {
    fn name(&self) -> &str {
        "StreamChecker"
    }

    fn check_health(&self, _document: &PdfDocument, data: &[u8]) -> Finding {
        let stream_count = data.windows(6).filter(|w| *w == b"stream").count();
        let mut metrics = HashMap::new();
        metrics.insert("stream_count".to_string(), stream_count as f64);

        Finding {
            check_name: "Stream Validation".to_string(),
            status: CheckStatus::Passed,
            details: format!("Found {} streams", stream_count),
            metrics,
        }
    }
}

/// Integrity health checker
pub struct IntegrityChecker;

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

impl IntegrityChecker {
    pub fn new() -> Self {
        Self
    }
}

impl HealthChecker for IntegrityChecker {
    fn name(&self) -> &str {
        "IntegrityChecker"
    }

    fn check_health(&self, _document: &PdfDocument, data: &[u8]) -> Finding {
        let has_eof = data.ends_with(b"%%EOF") || data.ends_with(b"%%EOF\n");
        let mut metrics = HashMap::new();
        metrics.insert("has_eof".to_string(), if has_eof { 1.0 } else { 0.0 });
        metrics.insert("file_size".to_string(), data.len() as f64);

        Finding {
            check_name: "Integrity Validation".to_string(),
            status: if has_eof {
                CheckStatus::Passed
            } else {
                CheckStatus::Warning
            },
            details: if has_eof {
                "File integrity appears intact".to_string()
            } else {
                "Missing EOF marker - file may be truncated".to_string()
            },
            metrics,
        }
    }
}

/// Quick health assessment function
pub fn quick_health_check(document: &PdfDocument, data: &[u8]) -> DocumentHealth {
    let diagnostics = DocumentDiagnostics::new(DiagnosticsConfig {
        deep_analysis: false,
        ..DiagnosticsConfig::default()
    });

    let report = diagnostics.analyze_health(document, data);
    report.overall_health
}