"""Tests for export control screening."""

from axiom.rag.ec_screening import screen_document, screen_batch


class TestECScreening:
    def test_clean_doc_allowed(self):
        r = screen_document("ORNL-4396.pdf", "MSRE operations experience report")
        assert r.allowed_community is True

    def test_ss_filename_flagged(self):
        r = screen_document("20071116 SS ML073180552.pdf", "Some content")
        assert r.allowed_community is False
        assert r.recommendation == "org"

    def test_ouo_content_flagged(self):
        r = screen_document("doc.pdf", "This document contains Official Use Only - Security Related Information")
        assert r.allowed_community is False

    def test_sunsi_flagged(self):
        r = screen_document("doc.pdf", "SUNSI Review Complete - document cleared")
        assert r.allowed_community is False

    def test_scientific_classified_not_flagged(self):
        """'classified as Class A experiment' is scientific, not security."""
        r = screen_document("doc.pdf", "experiments were classified as either Class A or Class B")
        assert r.allowed_community is True

    def test_regulatory_text_about_classification_not_flagged(self):
        """10 CFR discussing classification policy should pass."""
        r = screen_document("10CFR-Part50.pdf", "The Commission may classify information as restricted")
        assert r.allowed_community is True

    def test_org_corpus_not_screened(self):
        """Documents targeted for org tier skip screening."""
        r = screen_document("SECRET-doc.pdf", "TOP SECRET", target_corpus="rag-org")
        assert r.allowed_community is True

    def test_batch_screening(self):
        docs = [
            ("clean.pdf", "Normal reactor physics document"),
            ("SS flagged.pdf", "Some content"),
            ("also-clean.md", "More normal content"),
        ]
        allowed, flagged = screen_batch(docs)
        assert len(allowed) == 2
        assert len(flagged) == 1
        assert flagged[0].path == "SS flagged.pdf"
