//! OCR (Optical Character Recognition) integration for PDF
//!
//! Provides a pluggable OCR engine interface, Tesseract integration via
//! command-line, hOCR parsing, and searchable PDF generation.

use std::process::Command;
use std::sync::Mutex;

use thiserror::Error;

use crate::fitz::error::Error;
use crate::pdf::write;

// ============================================================================
// OcrError
// ============================================================================

/// OCR-specific error type
#[derive(Error, Debug)]
pub enum OcrError {
    #[error("Tesseract binary not found: {0}")]
    TesseractNotFound(String),
    #[error("Language not available: {0}")]
    LanguageNotAvailable(String),
    #[error("OCR execution failed: {0}")]
    ExecutionFailed(String),
    #[error("hOCR parsing failed: {0}")]
    HocrParseFailed(String),
    #[error("Invalid image dimensions")]
    InvalidImageDimensions,
    #[error("{0}")]
    Other(String),
}

impl From<OcrError> for Error {
    fn from(e: OcrError) -> Self {
        Error::Generic(e.to_string())
    }
}

pub type OcrResult<T> = std::result::Result<T, OcrError>;

// ============================================================================
// OcrWord
// ============================================================================

/// A single recognized word with bounding box and confidence
#[derive(Debug, Clone, PartialEq)]
pub struct OcrWord {
    /// Recognized text
    pub text: String,
    /// Bounding box: x, y, width, height (image coordinates: origin top-left)
    pub bbox: (f64, f64, f64, f64),
    /// Confidence score 0.0–100.0
    pub confidence: f64,
    /// Line number (0-based) for ordering
    pub line_number: u32,
}

impl OcrWord {
    /// Create a new OcrWord
    pub fn new(
        text: impl Into<String>,
        x: f64,
        y: f64,
        width: f64,
        height: f64,
        confidence: f64,
        line_number: u32,
    ) -> Self {
        Self {
            text: text.into(),
            bbox: (x, y, width, height),
            confidence,
            line_number,
        }
    }

    /// X coordinate (left edge)
    pub fn x(&self) -> f64 {
        self.bbox.0
    }

    /// Y coordinate (top edge in image coords)
    pub fn y(&self) -> f64 {
        self.bbox.1
    }

    /// Width
    pub fn width(&self) -> f64 {
        self.bbox.2
    }

    /// Height
    pub fn height(&self) -> f64 {
        self.bbox.3
    }
}

// ============================================================================
// OcrResult (recognition result, not Result<T,E>)
// ============================================================================

/// Result of OCR recognition on an image
#[derive(Debug, Clone)]
pub struct OcrRecognitionResult {
    /// Full recognized text
    pub text: String,
    /// Individual words with bounding boxes
    pub words: Vec<OcrWord>,
    /// Overall confidence 0.0–100.0
    pub confidence: f64,
}

impl OcrRecognitionResult {
    /// Create from words, computing full text and average confidence
    pub fn from_words(words: Vec<OcrWord>) -> Self {
        let text = words
            .iter()
            .map(|w| w.text.as_str())
            .collect::<Vec<_>>()
            .join(" ");
        let confidence = if words.is_empty() {
            0.0
        } else {
            words.iter().map(|w| w.confidence).sum::<f64>() / words.len() as f64
        };
        Self {
            text,
            words,
            confidence,
        }
    }

    /// Create empty result
    pub fn empty() -> Self {
        Self {
            text: String::new(),
            words: Vec::new(),
            confidence: 0.0,
        }
    }
}

// ============================================================================
// OcrConfig
// ============================================================================

/// Page segmentation mode (Tesseract PSM)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PageSegmentationMode {
    #[default]
    Auto = 3,
    SingleColumn = 4,
    SingleBlock = 6,
    SingleLine = 7,
    SingleWord = 8,
    SparseText = 11,
}

/// OCR engine configuration
#[derive(Debug, Clone)]
pub struct OcrConfig {
    /// Language code (e.g., "eng", "fra", "deu")
    pub language: String,
    /// Page segmentation mode
    pub page_segmentation_mode: PageSegmentationMode,
    /// DPI for image (affects scale)
    pub dpi: u32,
    /// Whitelist characters (e.g., "0123456789" for digits only)
    pub whitelist: Option<String>,
    /// Blacklist characters to exclude
    pub blacklist: Option<String>,
}

impl Default for OcrConfig {
    fn default() -> Self {
        Self {
            language: "eng".to_string(),
            page_segmentation_mode: PageSegmentationMode::Auto,
            dpi: 300,
            whitelist: None,
            blacklist: None,
        }
    }
}

impl OcrConfig {
    /// Validate configuration
    pub fn validate(&self) -> OcrResult<()> {
        if self.language.is_empty() {
            return Err(OcrError::Other("language cannot be empty".to_string()));
        }
        if self.dpi == 0 {
            return Err(OcrError::Other("dpi must be positive".to_string()));
        }
        Ok(())
    }
}

// ============================================================================
// OcrEngine trait
// ============================================================================

/// Pluggable OCR engine interface
pub trait OcrEngine {
    /// Recognize text in image bytes (raw RGB or grayscale, row-major)
    fn recognize(&self, image: &[u8], width: u32, height: u32) -> OcrResult<OcrRecognitionResult>;

    /// Supported language codes
    fn supported_languages(&self) -> Vec<String>;

    /// Set active language
    fn set_language(&mut self, lang: &str) -> OcrResult<()>;
}

// ============================================================================
// TesseractEngine
// ============================================================================

/// Tesseract OCR engine via command-line
pub struct TesseractEngine {
    /// Path to Tesseract data (e.g., /usr/share/tesseract-ocr/4.00/tessdata)
    data_path: Option<String>,
    /// Current language
    language: String,
    /// Page segmentation mode
    psm: PageSegmentationMode,
    /// Whitelist/blacklist
    whitelist: Option<String>,
    blacklist: Option<String>,
}

impl TesseractEngine {
    /// Create a new Tesseract engine
    /// `data_path`: optional path to tessdata directory
    pub fn new(data_path: Option<&str>) -> OcrResult<Self> {
        // Check that tesseract is available
        let output = Command::new("tesseract")
            .arg("--version")
            .output()
            .map_err(|e| OcrError::TesseractNotFound(e.to_string()))?;
        if !output.status.success() {
            return Err(OcrError::TesseractNotFound(
                String::from_utf8_lossy(&output.stderr).to_string(),
            ));
        }

        Ok(Self {
            data_path: data_path.map(String::from),
            language: "eng".to_string(),
            psm: PageSegmentationMode::Auto,
            whitelist: None,
            blacklist: None,
        })
    }

    /// Run tesseract and get hOCR output
    fn run_tesseract(&self, image: &[u8], width: u32, height: u32) -> OcrResult<String> {
        use std::io::Write;

        let mut cmd = Command::new("tesseract");
        cmd.stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .arg("stdin")
            .arg("stdout")
            .arg("-l")
            .arg(&self.language)
            .arg("hocr")
            .arg("--oem")
            .arg("3")
            .arg("-c")
            .arg(format!("tessedit_pageseg_mode={}", self.psm as u8));

        if let Some(ref path) = self.data_path {
            cmd.arg("--tessdata-dir").arg(path);
        }
        if let Some(ref wl) = self.whitelist {
            cmd.arg("-c").arg(format!("tessedit_char_whitelist={}", wl));
        }
        if let Some(ref bl) = self.blacklist {
            cmd.arg("-c").arg(format!("tessedit_char_blacklist={}", bl));
        }

        let mut child = cmd
            .spawn()
            .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
        if let Some(mut stdin) = child.stdin.take() {
            // Write PPM format (simplest for raw RGB)
            let header = format!("P6\n{} {}\n255\n", width, height);
            stdin
                .write_all(header.as_bytes())
                .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
            stdin
                .write_all(image)
                .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
        }
        let output = child
            .wait_with_output()
            .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
        if !output.status.success() {
            return Err(OcrError::ExecutionFailed(
                String::from_utf8_lossy(&output.stderr).to_string(),
            ));
        }
        String::from_utf8(output.stdout).map_err(|e| OcrError::ExecutionFailed(e.to_string()))
    }
}

impl OcrEngine for TesseractEngine {
    fn recognize(&self, image: &[u8], width: u32, height: u32) -> OcrResult<OcrRecognitionResult> {
        if width == 0 || height == 0 {
            return Err(OcrError::InvalidImageDimensions);
        }
        let expected = (width as usize) * (height as usize) * 3;
        if image.len() < expected {
            return Err(OcrError::InvalidImageDimensions);
        }

        let mut cmd = Command::new("tesseract");
        cmd.stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .arg("stdin")
            .arg("stdout")
            .arg("-l")
            .arg(&self.language)
            .arg("hocr")
            .arg("--oem")
            .arg("3")
            .arg("-c")
            .arg(format!("tessedit_pageseg_mode={}", self.psm as u8));

        if let Some(ref path) = self.data_path {
            cmd.arg("--tessdata-dir").arg(path);
        }
        if let Some(ref wl) = self.whitelist {
            cmd.arg("-c").arg(format!("tessedit_char_whitelist={}", wl));
        }
        if let Some(ref bl) = self.blacklist {
            cmd.arg("-c").arg(format!("tessedit_char_blacklist={}", bl));
        }

        let mut child = cmd
            .spawn()
            .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
        {
            let mut stdin = child
                .stdin
                .take()
                .ok_or_else(|| OcrError::ExecutionFailed("no stdin".to_string()))?;
            let header = format!("P6\n{} {}\n255\n", width, height);
            std::io::Write::write_all(&mut stdin, header.as_bytes())
                .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
            std::io::Write::write_all(&mut stdin, image)
                .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
        }
        let output = child
            .wait_with_output()
            .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
        if !output.status.success() {
            return Err(OcrError::ExecutionFailed(
                String::from_utf8_lossy(&output.stderr).to_string(),
            ));
        }
        let hocr = String::from_utf8(output.stdout)
            .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;

        match parse_hocr(&hocr) {
            Ok(words) if !words.is_empty() => Ok(OcrRecognitionResult::from_words(words)),
            Ok(_) => {
                // Fallback: plain text
                let mut cmd = Command::new("tesseract");
                cmd.stdin(std::process::Stdio::piped())
                    .stdout(std::process::Stdio::piped())
                    .stderr(std::process::Stdio::piped())
                    .arg("stdin")
                    .arg("stdout")
                    .arg("-l")
                    .arg(&self.language)
                    .arg("--oem")
                    .arg("3")
                    .arg("-c")
                    .arg(format!("tessedit_pageseg_mode={}", self.psm as u8));
                if let Some(ref path) = self.data_path {
                    cmd.arg("--tessdata-dir").arg(path);
                }
                let mut child = cmd
                    .spawn()
                    .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
                {
                    let mut stdin = child
                        .stdin
                        .take()
                        .ok_or_else(|| OcrError::ExecutionFailed("no stdin".to_string()))?;
                    let header = format!("P6\n{} {}\n255\n", width, height);
                    std::io::Write::write_all(&mut stdin, header.as_bytes())
                        .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
                    std::io::Write::write_all(&mut stdin, image)
                        .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
                }
                let out = child
                    .wait_with_output()
                    .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
                let text = String::from_utf8(out.stdout)
                    .map_err(|e| OcrError::ExecutionFailed(e.to_string()))?;
                Ok(OcrRecognitionResult {
                    text: text.trim().to_string(),
                    words: Vec::new(),
                    confidence: 0.0,
                })
            }
            Err(e) => Err(e),
        }
    }

    fn supported_languages(&self) -> Vec<String> {
        let output = Command::new("tesseract").arg("--list-langs").output();
        match output {
            Ok(o) if o.status.success() => {
                String::from_utf8_lossy(&o.stdout)
                    .lines()
                    .skip(1) // First line is "List of available languages:"
                    .map(|s| s.trim().to_string())
                    .filter(|s| !s.is_empty())
                    .collect()
            }
            _ => vec!["eng".to_string()],
        }
    }

    fn set_language(&mut self, lang: &str) -> OcrResult<()> {
        self.language = lang.to_string();
        Ok(())
    }
}

// Thread-safe wrapper for TesseractEngine (needed because set_language takes &mut self)
pub struct TesseractEngineSync(Mutex<TesseractEngine>);

impl TesseractEngineSync {
    pub fn new(data_path: Option<&str>) -> OcrResult<Self> {
        Ok(Self(Mutex::new(TesseractEngine::new(data_path)?)))
    }
}

impl OcrEngine for TesseractEngineSync {
    fn recognize(&self, image: &[u8], width: u32, height: u32) -> OcrResult<OcrRecognitionResult> {
        self.0
            .lock()
            .map_err(|_| OcrError::Other("lock poisoned".to_string()))?
            .recognize(image, width, height)
    }

    fn supported_languages(&self) -> Vec<String> {
        self.0
            .lock()
            .map(|e| e.supported_languages())
            .unwrap_or_else(|_| Vec::new())
    }

    fn set_language(&mut self, lang: &str) -> OcrResult<()> {
        self.0
            .lock()
            .map_err(|_| OcrError::Other("lock poisoned".to_string()))?
            .set_language(lang)
    }
}

// ============================================================================
// parse_hocr
// ============================================================================

/// Parse hOCR HTML output to extract words with bounding boxes
///
/// Expects `ocrx_word` spans with `title="bbox x0 y0 x1 y1; x_wconf N"`
pub fn parse_hocr(html: &str) -> OcrResult<Vec<OcrWord>> {
    let mut words = Vec::new();
    let mut line_number: u32 = 0;

    // Simple regex-less parsing: look for <span class="ocrx_word" title="bbox ...">
    let mut rest = html;
    while let Some(span_start) = rest.find("ocrx_word") {
        rest = &rest[span_start..];
        let title_start = match rest.find("title=\"") {
            Some(i) => i + 7,
            None => {
                rest = &rest[1..];
                continue;
            }
        };
        let title_end = match rest[title_start..].find('"') {
            Some(i) => title_start + i,
            None => {
                rest = &rest[1..];
                continue;
            }
        };
        let title = &rest[title_start..title_end];
        let (bbox, confidence) = parse_hocr_title(title)?;
        rest = &rest[title_end..];
        let content_start = match rest.find('>') {
            Some(i) => i + 1,
            None => {
                rest = &rest[1..];
                continue;
            }
        };
        let content_end = match rest[content_start..].find('<') {
            Some(i) => content_start + i,
            None => {
                rest = &rest[1..];
                continue;
            }
        };
        let text = html_entity_decode(&rest[content_start..content_end].trim());
        if !text.is_empty() {
            let (x0, y0, x1, y1) = bbox;
            let x = x0;
            let y = y0;
            let width = x1 - x0;
            let height = y1 - y0;
            words.push(OcrWord::new(
                text,
                x,
                y,
                width,
                height,
                confidence,
                line_number,
            ));
            line_number += 1;
        }
        rest = &rest[content_end..];
    }

    Ok(words)
}

fn parse_hocr_title(title: &str) -> OcrResult<((f64, f64, f64, f64), f64)> {
    let mut bbox = (0.0f64, 0.0f64, 0.0f64, 0.0f64);
    let mut confidence = 0.0f64;
    for part in title.split(';') {
        let part = part.trim();
        if part.starts_with("bbox ") {
            let coords: Vec<f64> = part[5..]
                .split_whitespace()
                .filter_map(|s| s.parse().ok())
                .collect();
            if coords.len() >= 4 {
                bbox = (coords[0], coords[1], coords[2], coords[3]);
            }
        } else if part.starts_with("x_wconf ") {
            if let Ok(c) = part[8..].trim().parse() {
                confidence = c;
            }
        }
    }
    Ok((bbox, confidence))
}

fn html_entity_decode(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut i = 0;
    let bytes = s.as_bytes();
    while i < bytes.len() {
        if bytes[i] == b'&' {
            if let Some(end) = bytes[i..].iter().position(|&b| b == b';') {
                let entity = &s[i + 1..i + end];
                match entity {
                    "lt" => out.push('<'),
                    "gt" => out.push('>'),
                    "amp" => out.push('&'),
                    "quot" => out.push('"'),
                    "apos" => out.push('\''),
                    _ if entity.starts_with("#x") || entity.starts_with("#X") => {
                        if let Ok(n) = u32::from_str_radix(&entity[2..], 16) {
                            if let Some(c) = char::from_u32(n) {
                                out.push(c);
                            }
                        }
                    }
                    _ if entity.starts_with('#') => {
                        if let Ok(n) = entity[1..].parse::<u32>() {
                            if let Some(c) = char::from_u32(n) {
                                out.push(c);
                            }
                        }
                    }
                    _ => {}
                }
                i += end + 1;
                continue;
            }
        }
        out.push(bytes[i] as char);
        i += 1;
    }
    out
}

// ============================================================================
// make_searchable_pdf
// ============================================================================

/// Generate PDF content stream for invisible text layer (render mode 3)
///
/// Positions text over the image so the PDF becomes searchable.
/// Image coordinates: origin top-left, y down. PDF coordinates: origin bottom-left, y up.
/// `image_height` is used to flip y.
pub fn make_searchable_pdf_content(
    result: &OcrRecognitionResult,
    image_width: u32,
    image_height: u32,
) -> Vec<u8> {
    if result.words.is_empty() {
        return Vec::new();
    }

    let h = image_height as f64;
    let mut out = Vec::new();

    // BT ... ET with Tr 3 (invisible)
    out.extend_from_slice(b"BT\n");
    out.extend_from_slice(b"3 Tr\n"); // Invisible text
    out.extend_from_slice(b"/F1 12 Tf\n"); // Default font, size 12

    for word in &result.words {
        let (x, y, width, height) = word.bbox;
        // PDF: y = 0 at bottom. Image: y = 0 at top. So pdf_y = h - y - height
        let pdf_y = h - y - height;
        let pdf_x = x;
        // Tm: a b c d e f = scale_x, skew_x, skew_y, scale_y, tx, ty
        let tm = format!("1 0 0 1 {} {} Tm\n", pdf_x, pdf_y);
        out.extend_from_slice(tm.as_bytes());
        let escaped = write::escape_pdf_string(word.text.as_bytes());
        out.extend_from_slice(b"(");
        out.extend_from_slice(&escaped);
        out.extend_from_slice(b") Tj\n");
    }

    out.extend_from_slice(b"ET\n");
    out
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_parse_hocr_single_word() {
        let html = r#"<span class="ocrx_word" id="word_1" title="bbox 10 20 50 40; x_wconf 95">Hello</span>"#;
        let words = parse_hocr(html).unwrap();
        assert_eq!(words.len(), 1);
        assert_eq!(words[0].text, "Hello");
        assert_eq!(words[0].bbox, (10.0, 20.0, 40.0, 20.0)); // x, y, width, height
        assert!((words[0].confidence - 95.0).abs() < 0.01);
    }

    #[test]
    fn test_parse_hocr_multiple_words() {
        let html = r#"
        <span class="ocrx_word" title="bbox 0 0 30 15; x_wconf 90">The</span>
        <span class="ocrx_word" title="bbox 35 0 80 15; x_wconf 88">quick</span>
        <span class="ocrx_word" title="bbox 85 0 120 15; x_wconf 92">brown</span>
        "#;
        let words = parse_hocr(html).unwrap();
        assert_eq!(words.len(), 3);
        assert_eq!(words[0].text, "The");
        assert_eq!(words[1].text, "quick");
        assert_eq!(words[2].text, "brown");
    }

    #[test]
    fn test_parse_hocr_html_entities() {
        let html =
            r#"<span class="ocrx_word" title="bbox 0 0 50 20; x_wconf 100">Hello&amp;World</span>"#;
        let words = parse_hocr(html).unwrap();
        assert_eq!(words.len(), 1);
        assert_eq!(words[0].text, "Hello&World");
    }

    #[test]
    fn test_parse_hocr_empty() {
        let words = parse_hocr("<div>no words</div>").unwrap();
        assert!(words.is_empty());
    }

    #[test]
    fn test_ocr_word_extraction() {
        let words = vec![
            OcrWord::new("test", 0.0, 0.0, 40.0, 15.0, 95.0, 0),
            OcrWord::new("word", 50.0, 0.0, 45.0, 15.0, 90.0, 1),
        ];
        let result = OcrRecognitionResult::from_words(words);
        assert_eq!(result.text, "test word");
        assert_eq!(result.words.len(), 2);
        assert!((result.confidence - 92.5).abs() < 0.01);
    }

    #[test]
    fn test_make_searchable_pdf_content() {
        let words = vec![
            OcrWord::new("Searchable", 10.0, 100.0, 80.0, 20.0, 95.0, 0),
            OcrWord::new("text", 100.0, 100.0, 30.0, 20.0, 90.0, 1),
        ];
        let result = OcrRecognitionResult::from_words(words);
        let content = make_searchable_pdf_content(&result, 200, 150);
        assert!(!content.is_empty());
        let s = String::from_utf8_lossy(&content);
        assert!(s.contains("BT"));
        assert!(s.contains("ET"));
        assert!(s.contains("3 Tr")); // Invisible render mode
        assert!(s.contains("Searchable"));
        assert!(s.contains("text"));
        // PDF y = 150 - 100 - 20 = 30
        assert!(s.contains("30"));
    }

    #[test]
    fn test_make_searchable_pdf_empty_words() {
        let result = OcrRecognitionResult::empty();
        let content = make_searchable_pdf_content(&result, 100, 100);
        assert!(content.is_empty());
    }

    #[test]
    fn test_ocr_config_default() {
        let config = OcrConfig::default();
        assert_eq!(config.language, "eng");
        assert_eq!(config.dpi, 300);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_ocr_config_validation_empty_language() {
        let mut config = OcrConfig::default();
        config.language = String::new();
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_ocr_config_validation_zero_dpi() {
        let mut config = OcrConfig::default();
        config.dpi = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_ocr_config_validation_ok() {
        let config = OcrConfig {
            language: "fra".to_string(),
            page_segmentation_mode: PageSegmentationMode::SingleColumn,
            dpi: 300,
            whitelist: Some("0123456789".to_string()),
            blacklist: None,
        };
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_ocr_error_display() {
        let e = OcrError::TesseractNotFound("not in PATH".to_string());
        assert!(e.to_string().contains("Tesseract"));
        assert!(e.to_string().contains("not in PATH"));
    }

    #[test]
    #[ignore] // Requires tesseract binary
    fn test_tesseract_engine_new() {
        let engine = TesseractEngine::new(None);
        if let Ok(e) = engine {
            let langs = e.supported_languages();
            assert!(!langs.is_empty());
            assert!(langs.contains(&"eng".to_string()));
        }
    }
}
