import unittest

from capstone_reference import capstone_inputs, run_capstone


class CapstoneTests(unittest.TestCase):
    def test_eligible_scope_and_labels_resolve(self):
        index, cases, total = capstone_inputs()
        self.assertEqual(total, 11)
        self.assertEqual(len(index.documents), 9)
        eligible = {document.id for document in index.documents}
        self.assertEqual(len(cases), 6)
        self.assertTrue(all(set(case["relevant_ids"]) <= eligible for case in cases))

    def test_context_budget_loss_is_reported_separately_from_retrieval(self):
        small = run_capstone(2, 80)
        larger = run_capstone(2, 400)
        self.assertEqual(small["retrieval"], larger["retrieval"])
        self.assertEqual(small["mean_packet_evidence_coverage"], 0)
        self.assertGreater(larger["mean_packet_evidence_coverage"], small["mean_packet_evidence_coverage"])
        self.assertFalse(larger["model_used"])

    def test_compound_case_retains_all_required_passages_when_budget_allows(self):
        report = run_capstone(2, 400)
        compound = next(row for row in report["context_rows"] if row["id"] == "cq-compound")
        self.assertEqual(set(compound["selected_ids"]), {"hours-v2", "study-room-v1"})
        self.assertEqual(compound["labeled_evidence_coverage"], 1)

    def test_runner_accepts_a_declared_selector_hook_without_changing_retrieval(self):
        def select_nothing(hits, max_quote_chars):
            return {"packet": {"evidence": []},
                    "budget_skipped_ids": [hit.document.id for hit in hits],
                    "conflicting_source_versions": {}}

        custom = run_capstone(2, 400, selector=select_nothing)
        baseline = run_capstone(2, 400)
        self.assertEqual(custom["retrieval"], baseline["retrieval"])
        self.assertEqual(custom["mean_packet_evidence_coverage"], 0)
        self.assertTrue(all(not row["selected_ids"] for row in custom["context_rows"]))


if __name__ == "__main__":
    unittest.main()