import unittest

from textbook.search_database import format_for_prompt, search, tokenize

INDEX = [
    {"name": "goodman_1_6_20", "id": "goodman-1-6-20", "source_label": "Corollary 1.6.20",
     "kind": "corollary", "status": "asserted",
     "statement": "\u2200 (p : \u2115) (as : List \u2124), Nat.Prime p \u2192 (p:\u2124) \u2223 List.prod as \u2192 \u2203 a \u2208 as, (p:\u2124) \u2223 a",
     "natural_language": "", "identifiers": ["Nat.Prime", "List.prod"], "requires": ["Nat.Prime"]},
    {"name": "goodman_1_6_24", "id": "goodman-1-6-24", "source_label": "Proposition 1.6.24",
     "kind": "proposition", "status": "asserted",
     "statement": "\u2200 (as : List \u2124), \u2203 d, IsGCDOfList d as",
     "natural_language": "", "identifiers": ["IsGCDOfList"], "requires": []},
    {"name": "goodman_1_6_1", "id": "goodman-1-6-1", "source_label": "Proposition 1.6.1",
     "kind": "proposition", "status": "asserted",
     "statement": "\u2200 (a : \u2124), a + 0 = a", "natural_language": "",
     "identifiers": [], "requires": []},
]


class SearchDatabaseTests(unittest.TestCase):
    def test_tokenize_keeps_qualified_names_and_drops_stopwords(self):
        tokens = tokenize('let x := Nat.Prime p')
        self.assertIn('Nat.Prime', tokens)
        self.assertNotIn('let', tokens)

    def test_exact_premise_ranks_first(self):
        hits = search(INDEX, '\u22a2 IsGCDOfList d as', limit=3)
        self.assertEqual(hits[0]['name'], 'goodman_1_6_24')

    def test_rare_identifier_outweighs_common_one(self):
        # Both premises mention List; only one mentions Nat.Prime, so the goal about
        # primes must prefer it rather than tying on the common token.
        hits = search(INDEX, '\u22a2 Nat.Prime p \u2192 (p:\u2124) \u2223 List.prod as', limit=3)
        self.assertEqual(hits[0]['name'], 'goodman_1_6_20')

    def test_irrelevant_premises_are_dropped_not_ranked_last(self):
        hits = search(INDEX, '\u22a2 IsGCDOfList d as', limit=5)
        self.assertNotIn('goodman_1_6_1', [h['name'] for h in hits])

    def test_proven_premise_is_preferred_over_an_asserted_one(self):
        index = [dict(INDEX[1]), dict(INDEX[1], name='proved_twin', id='proved-twin',
                                       status='proved', source_label='Proposition 1.6.24b')]
        hits = search(index, '\u22a2 IsGCDOfList d as', limit=2)
        self.assertEqual(hits[0]['name'], 'proved_twin')

    def test_limit_is_respected_and_empty_result_is_explained(self):
        self.assertEqual(len(search(INDEX, '\u22a2 List.prod as', limit=1)), 1)
        self.assertIn('no database premise', format_for_prompt([]))

    def test_format_for_prompt_marks_axioms_distinctly(self):
        text = format_for_prompt(search(INDEX, '\u22a2 IsGCDOfList d as', limit=1))
        self.assertIn('asserted from the book (axiom)', text)
        self.assertIn('goodman_1_6_24', text)


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