gaze-recognizers 0.9.1

use std::path::Path;
use std::sync::Arc;

use gaze_types::PiiClass;

use super::NerBackend;
use crate::ner::detector::NerDetector;
use crate::ner::error::NerRuntimeError;
use crate::ner::recognizer::NerRecognizer;
use crate::ner::types::{NerBackendKind, NerOptions, NerSpanResult};

struct TestSupportBackend;

struct ErrorBackend;

impl NerBackend for TestSupportBackend {
    fn detect(&self, input: &str) -> Result<Vec<NerSpanResult>, NerRuntimeError> {
        Ok(input
            .find("Alice Example")
            .or_else(|| input.find("Alice"))
            .map(|start| NerSpanResult {
                span: start
                    ..input[start..]
                        .find(|ch: char| !ch.is_alphabetic() && ch != ' ')
                        .map_or(input.len(), |end| start + end)
                        .min(start + "Alice Example".len()),
                class: PiiClass::Name,
                score: 0.40,
            })
            .into_iter()
            .collect())
    }
}

impl NerBackend for ErrorBackend {
    fn detect(&self, _input: &str) -> Result<Vec<NerSpanResult>, NerRuntimeError> {
        Err(NerRuntimeError::Inference(
            "forced test-support backend failure".to_string(),
        ))
    }
}

pub(crate) fn load_test_support_recognizer(
    model_dir: &Path,
    options: &NerOptions,
) -> Option<NerRecognizer> {
    let backend: Arc<dyn NerBackend> = match model_dir.file_name().and_then(|name| name.to_str()) {
        Some("__gaze_test_fixed_ner") => Arc::new(TestSupportBackend),
        Some("__gaze_test_error_ner") => Arc::new(ErrorBackend),
        _ => return None,
    };

    Some(NerRecognizer {
        detector: NerDetector {
            model_dir: model_dir.to_path_buf(),
            backend_kind: NerBackendKind::Ort,
            recognizer_version_id: "ner.test-support.v1".to_string(),
            locale: options.locale.clone(),
            threshold: options.threshold,
            backend,
        },
    })
}