package ru.pdguard;

import org.junit.jupiter.api.Test;
import ru.pdguard.config.SystemPolicy;
import ru.pdguard.core.PayloadStore;
import ru.pdguard.core.Pipeline;
import ru.pdguard.detect.NameCascade;
import ru.pdguard.detect.PdTypes;
import ru.pdguard.detect.RuleRegistry;
import ru.pdguard.detect.Span;
import ru.pdguard.mask.Masker;

import java.util.List;
import java.util.Optional;

import static org.junit.jupiter.api.Assertions.assertTrue;

/** LLAIM Legal NER: юридические реквизиты и документы, которых нет в общих моделях. */
class LegalNerTest {

    private List<Span> find(String text) {
        NameCascade cascade = new NameCascade(
                new NameCascade.EngineConfig(
                        "off", Optional.empty(),
                        "off", Optional.empty(),
                        "ru-legal-ner", Optional.of("models/ru-legal-ner")),
                16, 4);
        Pipeline p = new Pipeline(new RuleRegistry(), new Masker(), new PayloadStore(30), cascade);
        return p.findPersonalData(text, SystemPolicy.DEFAULT);
    }

    private boolean hasType(List<Span> spans, String type) {
        return spans.stream().anyMatch(s -> s.type().equals(type));
    }

    @Test
    void recognisesInn() {
        String text = "Договор между ООО «Ромашка», ИНН 7701234567, и Ивановым Иваном Ивановичем.";
        List<Span> spans = find(text);
        System.out.println("TEXT: " + text);
        for (Span s : spans) {
            System.out.println("  -> " + s.type() + " [" + text.substring(s.start(), s.end()) + "]");
        }
        assertTrue(hasType(spans, PdTypes.INN), "должно найти ИНН");
    }

    @Test
    void recognisesPassport() {
        String text = "Паспорт 4509 123456 выдан ОВД, СНИЛС 112-233-445 95.";
        List<Span> spans = find(text);
        System.out.println("TEXT: " + text);
        for (Span s : spans) {
            System.out.println("  -> " + s.type() + " [" + text.substring(s.start(), s.end()) + "]");
        }
        assertTrue(hasType(spans, PdTypes.PASSPORT), "должно найти паспорт");
    }
}