"""Exact-ranking teaching fixture, not an ANN engine or authorization service.
Run: python3 rag-filter-order.py
Tests: python3 rag-filter-order.py --test
No dependencies, network calls, or persistent writes.
"""
from dataclasses import dataclass
import json
import sys
import unittest


@dataclass(frozen=True)
class Document:
    id: str
    score: int
    allowed: bool
    relevant: bool


def retrieve(documents, k, mode, candidate_limit=None):
    if type(k) is not int or k < 1:
        raise ValueError("k must be a positive integer")
    if mode not in ("pre", "post"):
        raise ValueError("mode must be pre or post")
    if candidate_limit is None:
        candidate_limit = k
    if type(candidate_limit) is not int or candidate_limit < k:
        raise ValueError("candidate_limit must be an integer >= k")
    ranked = sorted(documents, key=lambda d: (-d.score, d.id))
    if len({d.id for d in ranked}) != len(ranked):
        raise ValueError("duplicate document IDs")
    if mode == "pre":
        return [d for d in ranked if d.allowed][:k]
    return [d for d in ranked[:candidate_limit] if d.allowed][:k]


def recall(documents, results):
    eligible = {d.id for d in documents if d.allowed and d.relevant}
    if not eligible:
        return None
    return len(eligible & {d.id for d in results}) / len(eligible)


def fixture():
    # Synthetic query scores; relevance labels are independent of the rank.
    return [Document(f"restricted-{i}", 100-i, False, True) for i in range(6)] + [
        Document("allowed-answer-a", 90, True, True),
        Document("allowed-answer-b", 89, True, True),
        Document("allowed-distraction", 88, True, False),
    ]


def demonstrate():
    docs = fixture()
    rows = []
    for mode, budget in [("post", 3), ("post", 6), ("post", 9), ("pre", 3)]:
        results = retrieve(docs, 3, mode, budget)
        rows.append({"mode": mode, "candidate_limit": budget,
                     "returned_ids": [d.id for d in results],
                     "authorised_relevant_recall": recall(docs, results)})
    return rows


class FilterOrderTests(unittest.TestCase):
    def test_post_filter_starves(self):
        self.assertEqual(retrieve(fixture(), 3, "post"), [])

    def test_pre_filter_finds_allowed(self):
        result = retrieve(fixture(), 3, "pre")
        self.assertEqual(len(result), 3)
        self.assertEqual(recall(fixture(), result), 1.0)

    def test_double_budget_still_starves(self):
        self.assertEqual(retrieve(fixture(), 3, "post", 6), [])

    def test_full_budget_recovers(self):
        self.assertEqual(retrieve(fixture(), 3, "post", 9), retrieve(fixture(), 3, "pre"))

    def test_both_modes_exclude_denied(self):
        for mode in ("pre", "post"):
            self.assertTrue(all(d.allowed for d in retrieve(fixture(), 3, mode, 9)))

    def test_all_denied(self):
        for mode in ("pre", "post"):
            self.assertEqual(retrieve(fixture()[:6], 3, mode), [])

    def test_no_eligible_relevance_is_undefined(self):
        self.assertIsNone(recall(fixture()[:6], []))

    def test_partial_recall(self):
        self.assertEqual(recall(fixture(), retrieve(fixture(), 1, "pre")), 0.5)

    def test_ties_are_deterministic(self):
        a = Document("a", 1, True, True)
        b = Document("b", 1, True, True)
        self.assertEqual(retrieve([b, a], 1, "pre"), [a])

    def test_reject_invalid_parameters(self):
        for k in (0, -1, True, 1.5):
            with self.assertRaises(ValueError):
                retrieve(fixture(), k, "pre")
        with self.assertRaises(ValueError):
            retrieve(fixture(), 3, "other")
        with self.assertRaises(ValueError):
            retrieve(fixture(), 3, "post", 2)

    def test_duplicate_ids_rejected(self):
        doc = fixture()[0]
        with self.assertRaises(ValueError):
            retrieve([doc, doc], 1, "pre")

    def test_empty_corpus(self):
        self.assertEqual(retrieve([], 3, "pre"), [])
        self.assertIsNone(recall([], []))


if __name__ == "__main__":
    if sys.argv[1:] == ["--test"]:
        unittest.main(argv=[sys.argv[0]])
    elif sys.argv[1:]:
        raise SystemExit("Usage: python3 rag-filter-order.py [--test]")
    else:
        print(json.dumps(demonstrate(), indent=2))