use std::path::Path;
use std::time::Instant;

use ocr::api::{Ocr, TextProcessor};
use ocr::core::config::{OcrConfig, PageSegMode, RecognitionEngine};
use ocr::lang::dictionary::Dictionary;

fn build_config(lang: &str, preprocess: bool, engine: &str, dict_correct: bool) -> OcrConfig {
    let mut config = OcrConfig::default();
    config.recognition.language = lang.to_string();
    config.recognition.confidence_threshold = 0.3;
    config.recognition.engine = match engine {
        "lstm" => RecognitionEngine::LSTM,
        "hybrid" => RecognitionEngine::Hybrid,
        _ => RecognitionEngine::PatternMatching,
    };
    config.recognition.enable_dictionary_correction = dict_correct;
    config.image_processing.enable_preprocessing = preprocess;
    config.image_processing.enable_binarization = preprocess;
    config.image_processing.enable_noise_reduction = preprocess;
    config.image_processing.enable_contrast_enhancement = preprocess;
    config.image_processing.enable_deskewing = preprocess;
    config.layout_analysis.page_seg_mode = PageSegMode::Auto;
    config
}

fn apply_dict_correction(text: &mut ocr::core::text::TextResult, lang: &str) {
    let mut dict = Dictionary::new();
    match lang {
        "en" => dict.load_words(&[
            "the", "this", "that", "and", "for", "are", "was", "but", "not", "you", "from", "they",
            "been", "with", "their", "would", "about", "which", "there", "could", "should",
            "people", "hello", "world", "ocr", "text",
        ]),
        "fr" => dict.load_words(&[
            "le", "de", "et", "être", "avoir", "pour", "pas", "dans", "une", "ils", "bonjour",
            "monde", "merci", "oui", "non", "travail", "amour", "vie",
        ]),
        "es" => dict.load_words(&[
            "el", "de", "que", "ser", "haber", "para", "con", "muy", "todo", "pero", "hola",
            "mundo", "gracias", "si", "no", "amor", "vida", "trabajo",
        ]),
        "de" => dict.load_words(&[
            "der", "die", "und", "sein", "haben", "für", "mit", "nicht", "auch", "auf", "hallo",
            "welt", "danke", "ja", "nein", "arbeit", "liebe", "leben",
        ]),
        "it" => dict.load_words(&[
            "il", "di", "che", "essere", "avere", "per", "con", "ma", "come", "non", "ciao",
            "mondo", "grazie", "sì", "no", "amore", "vita", "lavoro",
        ]),
        "pt" => dict.load_words(&[
            "o", "de", "que", "ser", "estar", "para", "com", "muito", "tudo", "mas", "olá",
            "mundo", "obrigado", "sim", "não", "amor", "vida", "trabalho",
        ]),
        "ru" => dict.load_words(&[
            "и",
            "в",
            "не",
            "на",
            "я",
            "быть",
            "он",
            "с",
            "что",
            "а",
            "привет",
            "мир",
            "спасибо",
            "да",
            "нет",
            "любовь",
            "жизнь",
            "работа",
        ]),
        _ => {}
    }
    for word in text.words.iter_mut() {
        let w = word.text.trim();
        if w.len() > 2 && !dict.contains(w) {
            let corrected = dict.correct_word(w);
            if corrected != w {
                word.text = corrected;
            }
        }
    }
    text.text = text
        .words
        .iter()
        .map(|w| w.text.as_str())
        .collect::<Vec<_>>()
        .join(" ");
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!(
            "Usage: {} <image-path> [engine] [lang] [preprocess] [dict-correct]",
            args.first().map(|s| s.as_str()).unwrap_or("ocr_demo")
        );
        eprintln!("  engine: pattern (default), lstm, hybrid");
        eprintln!("  lang: en (default), fr, es, de, it, pt, ru, zh, ja, ko, nl, pl, sv, da, fi, no, tr, el, hi, th, vi, ar, he, id, ms, uk, cs, hu, ro, bg");
        eprintln!("  preprocess: true (default), false");
        eprintln!("  dict-correct: true, false (default)");
        return Ok(());
    }

    let image_path = Path::new(&args[1]);
    let engine = args.get(2).map(|s| s.as_str()).unwrap_or("pattern");
    let lang = args.get(3).map(|s| s.as_str()).unwrap_or("en");
    let preprocess = args.get(4).map(|s| s == "true").unwrap_or(true);
    let dict_correct = args.get(5).map(|s| s == "true").unwrap_or(false);

    if !image_path.exists() {
        eprintln!("Error: image file not found: {:?}", image_path);
        return Ok(());
    }

    println!("OCR Demo");
    println!("========");
    println!("Image:     {:?}", image_path);
    println!("Engine:    {}", engine);
    println!("Language:  {}", lang);
    println!("Preprocess: {}", preprocess);
    println!("DictCorrect: {}", dict_correct);
    println!();

    let config = build_config(lang, preprocess, engine, dict_correct);
    let ocr = Ocr::with_config(config)?;
    let start = Instant::now();
    ocr.initialize().await?;
    println!("Initialized in {:?}", start.elapsed());

    let start = Instant::now();
    let mut result = ocr.recognize_text_from_file(image_path).await?;
    let elapsed = start.elapsed();
    println!("Recognized in {:?}", elapsed);

    if dict_correct {
        apply_dict_correction(&mut result, lang);
        println!("Dictionary correction applied");
    }

    println!();
    println!("Recognized text:");
    println!("{}", result.text);
    println!();
    println!("Confidence: {:.1}%", result.confidence * 100.0);
    println!("Words: {}", result.words.len());
    println!("Characters: {}", result.characters.len());
    println!();

    let filtered = TextProcessor::filter_by_confidence(&result, 0.5);
    if !filtered.text.is_empty() && filtered.text != result.text {
        println!("High-confidence text (>=50%):");
        println!("{}", filtered.text);
    }

    let stats = TextProcessor::get_text_statistics(&result);
    println!("Statistics:");
    println!("  Lines: {}", stats.line_count);
    println!("  Words: {}", stats.word_count);
    println!("  Characters: {}", stats.character_count);
    println!(
        "  Avg char confidence: {:.1}%",
        stats.avg_character_confidence * 100.0
    );
    println!(
        "  Avg word confidence: {:.1}%",
        stats.avg_word_confidence * 100.0
    );

    Ok(())
}
