df-ocr-switcher 0.1.0

//! Engine PaddleOCR PP-OCRv6 via ppocr-rs.
//!
//! Pipeline:
//!   DocOrientationClassifier → rotate_upright
//!   LayoutAnalyzer (PP-DocLayoutV3)
//!   detect_with_layout (OCR testo)
//!   [opt] TableStructureRecognizer (SLANeXt wired)
//!   [opt] FormulaRecognizer (PP-FormulaNet-plus-L)

use std::path::PathBuf;

use image::{imageops, RgbImage};
use ppocr_rs::{
    DocOrientation, DocOrientationClassifier, FormulaRecognizer,
    LayoutAnalyzer, LayoutBox, ModelHub, OcrLite, OcrOptions, Point,
    PpOcrVersion, PpStructureModel, SemanticClass,
    TableCellBox, TableStructureRecognizer, TextBlockWithLayout,
};

use crate::engine::{OcrBlock, OcrFormula, OcrLayoutBox, OcrPageResult, OcrTable, OcrWord};
use crate::error::{Error, Result};

use super::OcrEngine;

// ─── Config ──────────────────────────────────────────────────────────────────

/// Path espliciti per i modelli OCR (det/rec/dict).
pub struct OcrModelPaths {
    pub det:  PathBuf,
    pub rec:  PathBuf,
    pub dict: PathBuf,
}

/// Path espliciti per il modello table structure (SLANeXt wired/wireless o SLANet_plus).
pub struct TableModelPaths {
    /// Path al file `.onnx` (SLANeXt wired di default, oppure SLANet_plus).
    pub structure_onnx: PathBuf,
    /// Path al dizionario `table_structure_dict.txt`.
    pub structure_dict: PathBuf,
    /// Input size in pixel. `None` = default 512 (SLANeXt).
    /// Usare `Some(488)` per SLANet_plus.
    pub input_size: Option<u32>,
}

/// Configurazione per `PpOcrEngine`.
pub struct PpOcrEngineConfig {
    /// Versione PP-OCRv6: Tiny (default), Small, Medium.
    pub tier: PpOcrVersion,
    /// Path esplicito `inference.onnx` orientamento. Se `None` usa ModelHub.
    pub ori_model: Option<PathBuf>,
    /// Path esplicito per det/rec/dict OCR. Se `None` usa ModelHub.
    pub ocr_models: Option<OcrModelPaths>,
    /// Thread inferenza ONNX (default: 4).
    pub num_threads: usize,
    /// Path esplicito per SLANeXt table structure. Se `None` tenta ModelHub.
    pub table_models: Option<TableModelPaths>,
    /// Abilita table structure recognition. Default `true`.
    /// Se i modelli non sono in cache, la feature viene disabilitata silenziosamente.
    pub enable_tables: bool,
    /// Abilita il decoder formula (PP-FormulaNet → LaTeX). Default `false`:
    /// il decoder autoregressive è costoso (~3 s/formula su CPU).
    /// Attivare solo quando si vuole output LaTeX effettivo.
    pub enable_formula_decoder: bool,
}

impl Default for PpOcrEngineConfig {
    fn default() -> Self {
        Self {
            tier:                  PpOcrVersion::V6Tiny,
            ori_model:             None,
            ocr_models:            None,
            num_threads:           4,
            table_models:          None,
            enable_tables:         true,
            enable_formula_decoder: false,
        }
    }
}

// ─── Engine ──────────────────────────────────────────────────────────────────

pub struct PpOcrEngine {
    ocr:         OcrLite,
    ori:         Option<DocOrientationClassifier>,
    table_rec:   Option<TableStructureRecognizer>,
    formula_rec: Option<FormulaRecognizer>,
}

impl PpOcrEngine {
    /// Costruisce l'engine caricando i modelli.
    ///
    /// Table e formula recognizer sono opzionali: se i modelli non sono in
    /// cache e `fetch-models` non è attivo, la feature viene disabilitata
    /// con un warning su stderr ma l'engine rimane utilizzabile.
    pub fn new(cfg: PpOcrEngineConfig) -> Result<Self> {
        let hub = ModelHub::with_default_cache()?;

        // ── Modelli OCR ───────────────────────────────────────────────────
        let (det, rec, dict) = if let Some(p) = cfg.ocr_models {
            (p.det, p.rec, p.dict)
        } else {
            let paths = hub.ensure(cfg.tier)?;
            (paths.det_onnx, paths.rec_onnx, paths.dict_txt)
        };

        let mut ocr = OcrLite::new();
        ocr.init_models_no_angle(
            det.to_str().ok_or_else(|| Error::Other("det path non UTF-8".into()))?,
            rec.to_str().ok_or_else(|| Error::Other("rec path non UTF-8".into()))?,
            dict.to_str().ok_or_else(|| Error::Other("dict path non UTF-8".into()))?,
            cfg.num_threads,
        )?;

        // ── Modello orientamento (opzionale — fallisce silenziosamente) ──────
        let ori = if let Some(ori_path) = cfg.ori_model {
            match DocOrientationClassifier::from_path(&ori_path) {
                Ok(clf) => Some(clf),
                Err(e) => {
                    eprintln!("[df-ocr-switcher] orientamento disabilitato: {e}");
                    None
                }
            }
        } else {
            match hub.ensure_single(PpStructureModel::DocOrientation) {
                Ok(sp) => match DocOrientationClassifier::from_path(&sp.onnx) {
                    Ok(clf) => Some(clf),
                    Err(e)  => {
                        eprintln!("[df-ocr-switcher] orientamento disabilitato: {e}");
                        None
                    }
                },
                Err(_) => None, // modello non in cache → skip silenzioso
            }
        };

        // ── Table structure recognizer (opzionale) ────────────────────────
        let table_rec = if cfg.enable_tables {
            match load_table_recognizer(&hub, cfg.table_models) {
                Ok(rec) => Some(rec),
                Err(e)  => {
                    eprintln!("[df-ocr-switcher] table recognition disabilitato: {e}");
                    None
                }
            }
        } else {
            None
        };

        // ── Formula recognizer (opzionale, decoder disabilitato di default) ─
        let formula_rec = match load_formula_recognizer(&hub, cfg.enable_formula_decoder) {
            Ok(mut rec) => {
                rec.decoder_enabled = cfg.enable_formula_decoder;
                Some(rec)
            }
            Err(e) => {
                if cfg.enable_formula_decoder {
                    eprintln!("[df-ocr-switcher] formula recognition disabilitato: {e}");
                }
                None
            }
        };

        Ok(Self { ocr, ori, table_rec, formula_rec })
    }
}

impl OcrEngine for PpOcrEngine {
    fn process_page(&mut self, img: &RgbImage, layout: &mut LayoutAnalyzer) -> Result<OcrPageResult> {
        // ── Step 1: orientamento ──────────────────────────────────────────
        let (page_angle, upright) = if let Some(clf) = &self.ori {
            let (orient, _conf) = clf.classify(img)?;
            let degrees = orient.degrees();
            let rotated = rotate_upright(img, orient);
            (degrees, rotated)
        } else {
            (0u32, img.clone())
        };
        let (w, h) = upright.dimensions();

        // ── Step 2: layout-aware OCR ──────────────────────────────────────
        let opts = OcrOptions {
            return_word_box:     true,
            use_doc_orientation: false,
            ..OcrOptions::default()
        };
        let result = self.ocr.detect_with_layout(
            &upright, layout,
            10, 960, 0.6, 0.3, 1.6,
            false, false,
            opts,
        )?;

        // ── Step 3: table pipeline ────────────────────────────────────────
        let tables = if let Some(rec) = &self.table_rec {
            extract_tables(&upright, &result.layout_boxes, &result.blocks, rec)
        } else {
            vec![]
        };

        // ── Step 4: formula pipeline ──────────────────────────────────────
        let formulas = if let Some(rec) = &self.formula_rec {
            extract_formulas(&upright, &result.layout_boxes, rec)
        } else {
            vec![]
        };

        // ── Step 5: converti in OcrPageResult ────────────────────────────
        let layout_boxes = result.layout_boxes.iter().map(lb_to_ocr).collect();
        let (blocks, words) = blocks_and_words(&result.blocks);

        Ok(OcrPageResult {
            page_angle, page_width: w, page_height: h,
            layout_boxes, blocks, words, tables, formulas,
        })
    }
}

// ─── Table pipeline ───────────────────────────────────────────────────────────

fn extract_tables(
    img:    &RgbImage,
    lbs:    &[LayoutBox],
    blocks: &[TextBlockWithLayout],
    rec:    &TableStructureRecognizer,
) -> Vec<OcrTable> {
    lbs.iter()
        .enumerate()
        .filter(|(_, lb)| lb.class.semantic() == SemanticClass::Table)
        .filter_map(|(i, lb)| {
            let crop = crop_lb(img, lb);
            match rec.recognize(&crop) {
                Ok(structure) => {
                    let gfm = table_to_gfm(&structure.cell_boxes, blocks, lb);
                    Some(OcrTable { layout_idx: i as i32, gfm })
                }
                Err(e) => {
                    eprintln!("[df-ocr-switcher] SLANeXt errore su tabella {i}: {e}");
                    None
                }
            }
        })
        .collect()
}

/// Converte le cell_boxes + testo OCR full-page in GFM Markdown table.
///
/// Le coordinate delle cell_boxes sono relative al crop tabella
/// → vanno traslate in coordinate di pagina aggiungendo (lb.x, lb.y).
fn table_to_gfm(
    cells:  &[TableCellBox],
    blocks: &[TextBlockWithLayout],
    lb:     &LayoutBox,
) -> String {
    if cells.is_empty() {
        return String::new();
    }

    // Ordina per centroide Y poi X
    let mut indexed: Vec<(usize, f32, f32)> = cells.iter()
        .enumerate()
        .map(|(i, c)| (i, (c.x1 + c.x2) * 0.5, (c.y1 + c.y2) * 0.5))
        .collect();
    indexed.sort_by(|a, b| a.2.partial_cmp(&b.2)
        .unwrap_or(std::cmp::Ordering::Equal)
        .then(a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)));

    // Stima numero colonne: massima densità in Y (prima riga)
    let first_y = indexed.first().map(|e| e.2).unwrap_or(0.0);
    let n_cols  = indexed.iter().filter(|e| (e.2 - first_y).abs() < 20.0).count().max(1);

    let mut md = String::new();
    for (row_i, chunk) in indexed.chunks(n_cols).enumerate() {
        // Ordina celle della riga per X
        let mut row = chunk.to_vec();
        row.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

        md.push('|');
        for &(ci, _, _) in &row {
            let text = collect_cell_text(blocks, &cells[ci], lb.x, lb.y);
            md.push_str(&format!(" {} |", text.replace('|', "\\|")));
        }
        md.push('\n');

        if row_i == 0 {
            md.push('|');
            for _ in 0..row.len() { md.push_str("---|"); }
            md.push('\n');
        }
    }
    md
}

/// Trova il testo dei blocchi OCR il cui centroide cade dentro la cella.
/// `lb_x / lb_y` = offset del crop tabella in coordinate pagina.
fn collect_cell_text(
    blocks: &[TextBlockWithLayout],
    cell:   &TableCellBox,
    lb_x:   u32,
    lb_y:   u32,
) -> String {
    let cx1 = lb_x as f32 + cell.x1;
    let cy1 = lb_y as f32 + cell.y1;
    let cx2 = lb_x as f32 + cell.x2;
    let cy2 = lb_y as f32 + cell.y2;

    let mut parts: Vec<(u32, &str)> = blocks.iter()
        .filter(|b| {
            let bx = b.centroid_x as f32;
            let by = b.centroid_y as f32;
            bx >= cx1 && bx <= cx2 && by >= cy1 && by <= cy2
        })
        .map(|b| (b.centroid_y, b.block.text.trim()))
        .collect();
    parts.sort_by_key(|(y, _)| *y);
    parts.iter()
        .filter(|(_, t)| !t.is_empty())
        .map(|(_, t)| *t)
        .collect::<Vec<_>>()
        .join(" ")
}

// ─── Formula pipeline ─────────────────────────────────────────────────────────

fn extract_formulas(
    img:  &RgbImage,
    lbs:  &[LayoutBox],
    rec:  &FormulaRecognizer,
) -> Vec<OcrFormula> {
    lbs.iter()
        .enumerate()
        .filter(|(_, lb)| lb.class.semantic() == SemanticClass::Equation)
        .filter_map(|(i, lb)| {
            let crop = crop_lb(img, lb);
            match rec.recognize(&crop) {
                Ok(fr) => Some(OcrFormula { layout_idx: i as i32, latex: fr.latex }),
                Err(e) => {
                    eprintln!("[df-ocr-switcher] formula rec errore su regione {i}: {e}");
                    None
                }
            }
        })
        .collect()
}

// ─── Model loaders ───────────────────────────────────────────────────────────

fn load_table_recognizer(
    hub:   &ModelHub,
    paths: Option<TableModelPaths>,
) -> std::result::Result<TableStructureRecognizer, Box<dyn std::error::Error>> {
    let (onnx, dict, input_size) = if let Some(p) = paths {
        (p.structure_onnx, p.structure_dict, p.input_size)
    } else {
        let sp = hub.ensure_single(PpStructureModel::TableStructureWired)?;
        let onnx = sp.onnx;
        let dict = sp.dict_txt.ok_or("SLANeXt: dict_txt mancante")?;
        (onnx, dict, None)
    };
    let rec = TableStructureRecognizer::from_path_with_dict(&onnx, Some(&dict))?;
    Ok(if let Some(sz) = input_size { rec.with_input_size(sz) } else { rec })
}

fn load_formula_recognizer(
    hub:     &ModelHub,
    enabled: bool,
) -> std::result::Result<FormulaRecognizer, Box<dyn std::error::Error>> {
    if !enabled {
        // Carica in modalità stub (decoder_enabled = false): nessuna inferenza
        // ma struttura inizializzata. Se i modelli non sono in cache → skip.
        let sp = hub.ensure_single(PpStructureModel::FormulaRec)?;
        let tok = sp.tokenizer_json.as_deref();
        Ok(FormulaRecognizer::from_paths(&sp.onnx, tok)?)
    } else {
        let sp = hub.ensure_single(PpStructureModel::FormulaRec)?;
        let tok = sp.tokenizer_json.as_deref();
        Ok(FormulaRecognizer::from_paths(&sp.onnx, tok)?)
    }
}

// ─── Helper ──────────────────────────────────────────────────────────────────

/// Crop di un LayoutBox dall'immagine.
fn crop_lb(img: &RgbImage, lb: &LayoutBox) -> RgbImage {
    let x = lb.x.min(img.width().saturating_sub(1));
    let y = lb.y.min(img.height().saturating_sub(1));
    let w = lb.w.min(img.width().saturating_sub(x));
    let h = lb.h.min(img.height().saturating_sub(y));
    imageops::crop_imm(img, x, y, w, h).to_image()
}

fn rotate_upright(img: &RgbImage, orient: DocOrientation) -> RgbImage {
    match orient {
        DocOrientation::Deg0   => img.clone(),
        DocOrientation::Deg90  => imageops::rotate270(img),
        DocOrientation::Deg180 => imageops::rotate180(img),
        DocOrientation::Deg270 => imageops::rotate90(img),
    }
}

fn aabb(pts: &[Point]) -> (u32, u32, u32, u32) {
    let x1 = pts.iter().map(|p| p.x).min().unwrap_or(0);
    let y1 = pts.iter().map(|p| p.y).min().unwrap_or(0);
    let x2 = pts.iter().map(|p| p.x).max().unwrap_or(0);
    let y2 = pts.iter().map(|p| p.y).max().unwrap_or(0);
    (x1, y1, x2, y2)
}

fn lb_to_ocr(lb: &LayoutBox) -> OcrLayoutBox {
    OcrLayoutBox {
        class_name:    format!("{:?}", lb.class),
        semantic:      lb.class.semantic(),
        x1: lb.xmin(), y1: lb.ymin(),
        x2: lb.xmax(), y2: lb.ymax(),
        reading_order: lb.reading_order,
    }
}

fn blocks_and_words(
    src: &[TextBlockWithLayout],
) -> (Vec<OcrBlock>, Vec<OcrWord>) {
    let mut blocks = Vec::with_capacity(src.len());
    let mut words  = Vec::new();

    for blk in src {
        let layout_idx = blk.layout_index.map(|i| i as i32).unwrap_or(-1);
        let (x1, y1, x2, y2) = aabb(&blk.block.box_points);

        blocks.push(OcrBlock {
            text:       blk.block.text.clone(),
            x1, y1, x2, y2,
            confidence: blk.block.text_score,
            layout_idx,
        });

        if blk.block.words.is_empty() {
            words.push(OcrWord {
                text: blk.block.text.clone(),
                x1, y1, x2, y2,
                confidence: blk.block.text_score,
                layout_idx,
            });
        } else {
            for w in &blk.block.words {
                let (wx1, wy1, wx2, wy2) = aabb(&w.box_points);
                words.push(OcrWord {
                    text: w.text.clone(),
                    x1: wx1, y1: wy1, x2: wx2, y2: wy2,
                    confidence: w.score,
                    layout_idx,
                });
            }
        }
    }

    (blocks, words)
}