//! Обёртка над предобученной ONNX-моделью (BERT-подобный классификатор
//! токсичности текста), выполняющая инференс без Python-рантайма — только
//! `tract-onnx` (чистый Rust ONNX-runtime) и `tokenizers` для препроцессинга.
//! Это позволяет держать сервис лёгким и без внешних зависимостей уровня ОС.

use anyhow::{anyhow, Result};
use tokenizers::Tokenizer;
use tract_onnx::prelude::*;

pub struct ToxicityClassifier {
    model: SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
    tokenizer: Tokenizer,
}

impl ToxicityClassifier {
    /// Загружает токенизатор и ONNX-модель с диска и компилирует граф модели
    /// в исполняемый план (`into_optimized().into_runnable()`).
    ///
    /// Вызывается один раз при старте сервиса (см. `main.rs`) — сама загрузка
    /// и оптимизация графа занимают заметное время, поэтому переиспользуем
    /// результат на все последующие запросы через `AppState`.
    pub fn new(model_path: &str, tokenizer_path: &str) -> Result<Self> {
        let tokenizer = Tokenizer::from_file(tokenizer_path)
            .map_err(|e| anyhow!("Tokenizer load error: {}", e))?;

        let model = tract_onnx::onnx()
            .model_for_path(model_path)?
            .into_optimized()?
            .into_runnable()?;

        Ok(Self { model, tokenizer })
    }

    /// Определяет, является ли текст токсичным.
    ///
    /// Пайплайн: токенизация -> три входных тензора (`ids`/`attention_mask`/
    /// `token_type_ids`) формы `[1, seq_len]` -> прогон через модель ->
    /// логиты -> сигмоида на логите токсичного класса -> порог 0.85.
    ///
    /// Стратегия обработки ошибок — намеренно "fail-open": любая ошибка на
    /// любом шаге (токенизация, инференс, неожиданная форма выходного
    /// тензора) приводит к `return false`, т.е. текст считается нетоксичным
    /// и пропускается дальше. Это осознанный выбор в пользу доступности
    /// сервиса: лучше изредка пропустить нежелательный текст при сбое модели,
    /// чем уронить или заблокировать всю гостевую книгу из-за проблем с
    /// инференсом (повреждённые данные, несовместимая версия рантайма и т.п.).
    pub fn is_toxic(&self, text: &str) -> bool {
        let encoding = match self.tokenizer.encode(text, true) {
            Ok(e) => e,
            Err(_) => return false,
        };

        let ids: Vec<i64> = encoding.get_ids().iter().map(|&x| x as i64).collect();
        let mask: Vec<i64> = encoding
            .get_attention_mask()
            .iter()
            .map(|&x| x as i64)
            .collect();
        let type_ids: Vec<i64> = encoding.get_type_ids().iter().map(|&x| x as i64).collect();

        let seq_len = ids.len();

        let ids_tensor = tensor1(&ids)
            .into_shape(&[1, seq_len])
            .expect("Shape error");
        let mask_tensor = tensor1(&mask)
            .into_shape(&[1, seq_len])
            .expect("Shape error");
        let type_tensor = tensor1(&type_ids)
            .into_shape(&[1, seq_len])
            .expect("Shape error");

        let result = self.model.run(tvec![
            ids_tensor.into(),
            mask_tensor.into(),
            type_tensor.into()
        ]);

        let outputs = match result {
            Ok(out) => out,
            Err(e) => {
                println!("💥 [MODEL ERR]: {:?}", e);
                return false;
            }
        };

        // --- БРОНЕБОЙНОЕ РЕШЕНИЕ (Никаких паник ndarray) ---
        // to_array_view может вернуть Err, если фактическая форма/тип выходного
        // тензора отличается от ожидаемого (например, из-за другой версии
        // модели) — в этом случае, как и везде здесь, просто fail-open.
        let view = match outputs[0].to_array_view::<f32>() {
            Ok(v) => v,
            Err(_) => return false,
        };

        // Превращаем любую матрицу (хоть 2D, хоть 3D) в плоский список чисел
        let logits: Vec<f32> = view.iter().copied().collect();

        // Теперь logits это просто одномерный массив, например [-2.5, 4.8].
        // Модель бинарная: индекс 0 — логит класса "нетоксично", индекс 1 —
        // логит класса "токсично". Если модель почему-то вернула один логит
        // (не two-class head), используем его напрямую как запасной вариант.
        let toxic_logit = if logits.len() >= 2 {
            logits[1] // Вероятность мата
        } else if logits.len() == 1 {
            logits[0]
        } else {
            return false;
        };

        // Применяем сигмоиду
        let prob = 1.0 / (1.0 + (-toxic_logit).exp());

        println!(
            "🤖 [AI DEBUG] Text: '{}' | Logits: {:?} | Prob: {:.4}",
            text, logits, prob
        );

        prob > 0.85 // Порог 85%
    }
}
