"""Export control screening for RAG ingest.

Checks document text and filenames for security markings before
allowing ingestion into the community (public) corpus. Documents
with security markers are flagged for human review or automatically
routed to the org (restricted) corpus.

Runs during ingest — never allows EC content into community tier.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass

log = logging.getLogger(__name__)

# Filename patterns that indicate security-sensitive content
_FILENAME_MARKERS = [
    r"\bSS\b",  # Security Sensitive
    r"\bsec\b",  # Security
    r"\bOUO\b",  # Official Use Only
    r"\bSGI\b",  # Safeguards Information
    r"\bSUNSI\b",  # Sensitive Unclassified Non-Safeguards Information
    r"CONFIDENTIAL",
    r"RESTRICTED",
]

# Content patterns that indicate actual security markings
# (not just regulatory text discussing classification)
_CONTENT_MARKERS = [
    (r"Official Use Only[\s\-]+Security Related", "OUO-Security"),
    (r"OFFICIAL USE ONLY[\s\-]+SECURITY", "OUO-Security"),
    (r"Safeguards Information[\s\-]+Modified Handling", "SGI"),
    (r"SUNSI[\s\-]Review Complete", "SUNSI"),
    (r"Export Controlled Information", "EC"),
    (r"UCNI[\s\-]", "UCNI"),
    (r"Naval Nuclear Propulsion Information", "NNPI"),
    (r"This document contains.*Official Use Only", "OUO-Document"),
    (r"DISTRIBUTION LIMITED", "Distribution-Limited"),
    (r"NOT FOR PUBLIC RELEASE", "Not-Public"),
]


@dataclass
class ScreeningResult:
    """Result of EC screening for a document."""

    path: str
    allowed_community: bool = True
    markers_found: list[str] | None = None
    recommendation: str = ""  # "community", "org", "review", "reject"


def screen_document(
    path: str,
    text: str,
    target_corpus: str = "rag-community",
) -> ScreeningResult:
    """Screen a document for export control markers before ingest.

    Args:
        path: File path (checked for filename markers)
        text: Document text (checked for content markers)
        target_corpus: Intended corpus ("rag-community" triggers screening)

    Returns:
        ScreeningResult with recommendation
    """
    # Only screen documents targeted for community (public) tier
    if target_corpus != "rag-community":
        return ScreeningResult(path=path, allowed_community=True, recommendation=target_corpus)

    markers = []

    # Check filename
    for pattern in _FILENAME_MARKERS:
        if re.search(pattern, path):
            markers.append(f"filename:{pattern}")

    # Check content (first 5000 chars — markings are typically in headers/footers)
    header = text[:5000] if text else ""
    footer = text[-2000:] if text and len(text) > 2000 else ""
    check_text = header + "\n" + footer

    for pattern, label in _CONTENT_MARKERS:
        if re.search(pattern, check_text, re.IGNORECASE):
            markers.append(f"content:{label}")

    if markers:
        log.warning("EC screening flagged %s: %s", path, markers)
        return ScreeningResult(
            path=path,
            allowed_community=False,
            markers_found=markers,
            recommendation="org",  # Route to restricted tier
        )

    return ScreeningResult(path=path, allowed_community=True, recommendation="community")


def screen_batch(
    documents: list[tuple[str, str]],
    target_corpus: str = "rag-community",
) -> tuple[list[tuple[str, str]], list[ScreeningResult]]:
    """Screen a batch of documents. Returns (allowed, flagged).

    Args:
        documents: List of (path, text) tuples
        target_corpus: Intended corpus

    Returns:
        (allowed_docs, flagged_results)
    """
    allowed = []
    flagged = []

    for path, text in documents:
        result = screen_document(path, text, target_corpus)
        if result.allowed_community:
            allowed.append((path, text))
        else:
            flagged.append(result)

    if flagged:
        log.warning(
            "EC screening: %d/%d documents flagged, routed to org tier",
            len(flagged),
            len(documents),
        )

    return allowed, flagged
