import math
import unittest

from rag_toolkit import *


class TestChunking(unittest.TestCase):
    TEXT = " ".join(f"w{i}" for i in range(50))

    def test_fixed_no_overlap_covers_everything_once(self):
        cs = chunk_fixed("d", self.TEXT, size=20)
        self.assertEqual([len(c.text.split()) for c in cs], [20, 20, 10])
        self.assertEqual(" ".join(c.text for c in cs), self.TEXT)

    def test_overlap_repeats_words_between_neighbors(self):
        cs = chunk_fixed("d", self.TEXT, size=20, overlap=5)
        self.assertEqual(cs[0].text.split()[-5:], cs[1].text.split()[:5])
        self.assertEqual(cs[-1].text.split()[-1], "w49")

    def test_bad_chunk_arguments(self):
        with self.assertRaises(ValueError):
            chunk_fixed("d", self.TEXT, size=10, overlap=10)

    def test_sentence_chunks_never_cut_a_sentence(self):
        text = "One two three. Four five six seven. Eight nine."
        cs = chunk_sentences("d", text, max_words=6)
        self.assertEqual([c.text for c in cs], ["One two three.", "Four five six seven. Eight nine."])

    def test_context_prefix_is_indexed_but_text_is_unchanged(self):
        c = Chunk("d#0", "d", "The company grew 3%.")
        out = add_context([c], {"d": "ACME Q2 report"})[0]
        self.assertEqual(out.text, c.text)
        self.assertIn("ACME Q2 report", out.indexed_text)

    def test_contextual_prefix_lets_bm25_find_a_chunk_with_no_name_in_it(self):
        chunks = [Chunk("a#0", "a", "Revenue grew by three percent."), Chunk("b#0", "b", "Revenue fell by two percent.")]
        plain = BM25(chunks).search("ACME revenue grew")
        ctx = BM25(add_context(chunks, {"a": "ACME report", "b": "Zeta report"})).search("ACME revenue grew")
        self.assertEqual(ctx[0][0], "a#0")
        self.assertEqual(plain[0][0], "a#0")           # 'grew' already helped here...
        self.assertGreater(ctx[0][1], plain[0][1])      # ...but the name raises the score


class TestBM25(unittest.TestCase):
    def setUp(self):
        self.chunks = [Chunk("a", "d", "error ERR-4042 means the key file is missing"),
                       Chunk("b", "d", "the cat sat on the mat"),
                       Chunk("c", "d", "the dog and the cat play in the park every day")]
        self.bm = BM25(self.chunks)

    def test_exact_rare_term_wins(self):
        self.assertEqual(self.bm.search("ERR-4042")[0][0], "a")

    def test_no_overlap_means_no_results(self):
        self.assertEqual(self.bm.search("zebra"), [])

    def test_rarer_terms_have_higher_idf(self):
        self.assertGreater(self.bm.idf("4042"), self.bm.idf("cat"))
        self.assertGreater(self.bm.idf("cat"), 0)

    def test_repeating_a_term_has_diminishing_returns(self):
        s = [BM25([Chunk("x", "d", " ".join(["apple"] * n) + " filler" * 5), Chunk("y", "d", "other words here")]).score("apple", 0)
             for n in (1, 2, 4, 8)]
        gains = [b - a for a, b in zip(s, s[1:])]
        self.assertTrue(all(g > 0 for g in gains))
        self.assertGreater(gains[0], gains[2])
        self.assertLess(s[-1], 3 * s[0])                 # saturates: 8x the count is far less than 8x the score

    def test_shorter_document_scores_higher_for_same_count(self):
        bm = BM25([Chunk("s", "d", "apple pie"), Chunk("l", "d", "apple " + "filler " * 40)])
        self.assertGreater(bm.score("apple", 0), bm.score("apple", 1))


class TestDenseAndFusion(unittest.TestCase):
    def test_embeddings_are_unit_length_and_deterministic(self):
        e = ToyEmbedder().fit(["alpha beta", "gamma"])
        v1, v2 = e.embed("alpha beta"), e.embed("alpha beta")
        self.assertEqual(v1, v2)
        self.assertAlmostEqual(math.sqrt(sum(x * x for x in v1)), 1.0)

    def test_synonym_table_lets_dense_match_a_paraphrase_that_bm25_misses(self):
        chunks = [Chunk("a", "d", "Each team gets 50 dollars of credit."), Chunk("b", "d", "The dining hall opens at 7 am.")]
        q = "how much cash for a group"
        self.assertEqual(BM25(chunks).search(q), [])
        self.assertEqual(DenseIndex(chunks).search(q)[0][0], "a")

    def test_rrf_formula(self):
        fused = dict(rrf([["a", "b"], ["b", "c"]], k=60))
        self.assertAlmostEqual(fused["a"], 1 / 61)
        self.assertAlmostEqual(fused["b"], 1 / 62 + 1 / 61)
        self.assertAlmostEqual(fused["c"], 1 / 62)

    def test_rrf_rewards_agreement_between_lists(self):
        top = rrf([["x", "a", "b"], ["y", "a", "c"]])[0][0]
        self.assertEqual(top, "a")

    def test_rrf_handles_one_list_and_empty_input(self):
        self.assertEqual([d for d, _ in rrf([["a", "b"]])], ["a", "b"])
        self.assertEqual(rrf([]), [])


class TestLateInteraction(unittest.TestCase):
    def test_maxsim_of_identical_vectors_equals_number_of_query_tokens(self):
        e = ToyEmbedder()
        v = e.embed_tokens("alpha beta gamma")
        self.assertAlmostEqual(maxsim(v, v), 3.0)

    def test_maxsim_prefers_the_passage_that_covers_more_query_tokens(self):
        e = ToyEmbedder()
        q = e.embed_tokens("alpha beta gamma")
        both = e.embed_tokens("alpha beta something")
        one = e.embed_tokens("alpha delta epsilon")
        self.assertGreater(maxsim(q, both), maxsim(q, one))

    def test_rerank_returns_the_same_ids_reordered(self):
        pipe = Pipeline(build_chunks())
        ids = pipe.search("Which session has the shortest queue?", 6, "hybrid")
        out = rerank_late_interaction("Which session has the shortest queue?", pipe.by_id, ids, pipe.dense.emb)
        self.assertEqual(sorted(i for i, _ in out), sorted(ids))
        self.assertEqual(out[0][0], "office#0")


class TestMetrics(unittest.TestCase):
    def test_hit_recall_mrr(self):
        ranked, rel = ["a", "b", "c", "d"], ["c", "z"]
        self.assertEqual(hit_at_k(ranked, rel, 2), 0.0)
        self.assertEqual(hit_at_k(ranked, rel, 3), 1.0)
        self.assertEqual(recall_at_k(ranked, rel, 4), 0.5)
        self.assertAlmostEqual(mrr(ranked, rel), 1 / 3)
        self.assertEqual(mrr(ranked, ["q"]), 0.0)
        self.assertEqual(mrr(ranked, rel, k=2), 0.0)

    def test_ndcg(self):
        self.assertAlmostEqual(ndcg_at_k(["a", "b"], ["a"], 2), 1.0)
        self.assertAlmostEqual(ndcg_at_k(["x", "a"], ["a"], 2), 1 / math.log2(3))
        self.assertAlmostEqual(ndcg_at_k(["a", "b", "c"], ["a", "c"], 3), (1 + 1 / 2) / (1 + 1 / math.log2(3)))
        self.assertEqual(ndcg_at_k(["a"], [], 3), 0.0)

    def test_relevant_ids_finds_chunks_containing_the_answer(self):
        cs = [Chunk("a", "d", "The answer is 42."), Chunk("b", "d", "nothing")]
        self.assertEqual(relevant_ids(cs, "answer is 42"), ["a"])


class TestGenerator(unittest.TestCase):
    def setUp(self):
        self.pipe = Pipeline(build_chunks())

    def ask(self, q):
        retrieved = [self.pipe.by_id[i] for i in self.pipe.search(q, 3, "hybrid+rerank")]
        return mock_llm_answer(q, retrieved), retrieved

    def test_answers_with_a_citation(self):
        a, retrieved = self.ask("What does ERR-4042 mean?")
        self.assertFalse(a.abstained)
        self.assertIn("API key file is missing", a.text)
        self.assertEqual(a.citations, ["setup#1"])
        self.assertTrue(check_citations(a, retrieved))

    def test_abstains_when_sources_do_not_help(self):
        a, _ = self.ask(UNANSWERABLE)
        self.assertTrue(a.abstained)
        self.assertEqual(a.text, NO_ANSWER)
        self.assertEqual(a.citations, [])

    def test_abstains_with_no_sources(self):
        self.assertTrue(mock_llm_answer("anything at all", []).abstained)

    def test_a_fake_citation_is_caught(self):
        _, retrieved = self.ask("What does ERR-4042 mean?")
        self.assertFalse(check_citations(Answer("x [zzz#9]", ["zzz#9"]), retrieved))

    def test_every_lab_question_is_answered_from_the_right_chunk(self):
        for q, needle in QUERIES:
            a, _ = self.ask(q)
            self.assertFalse(a.abstained, q)
            self.assertIn(needle.lower(), a.text.lower(), q)


class TestPipelineAndAgent(unittest.TestCase):
    def setUp(self):
        self.pipe = Pipeline(build_chunks())

    def test_hybrid_is_never_worse_than_the_worst_single_retriever_on_the_lab_set(self):
        scores = {m: evaluate(self.pipe, m)[2] for m in ("bm25", "dense", "hybrid")}
        self.assertGreaterEqual(scores["hybrid"], min(scores["bm25"], scores["dense"]))
        self.assertEqual(scores["hybrid"], 1.0)

    def test_fixed_chunks_can_cut_an_answer_in_half(self):
        cut = Pipeline(build_chunks("fixed", size=12))
        whole = Pipeline(build_chunks("sentence", size=12))
        self.assertLess(evaluate(cut, "hybrid")[1], evaluate(whole, "hybrid")[1])

    def test_needs_retrieval(self):
        self.assertFalse(needs_retrieval("hello"))
        self.assertFalse(needs_retrieval("Thanks!"))
        self.assertTrue(needs_retrieval("hello, when is the milestone due?"))

    def test_agent_skips_retrieval_for_small_talk(self):
        calls = []
        ans, trace = agentic_answer("hi", lambda q, k: calls.append(q) or [], self.pipe.by_id)
        self.assertEqual(calls, [])
        self.assertIn("skip", trace[0])

    def test_agent_stops_with_abstain_and_does_not_loop_forever(self):
        calls = []
        def search(q, k):
            calls.append(q)
            return self.pipe.search(q, k)
        ans, trace = agentic_answer(UNANSWERABLE, search, self.pipe.by_id, max_rounds=3)
        self.assertTrue(ans.abstained)
        self.assertLessEqual(len(calls), 3)
        self.assertIn("stop", trace[-1])

    def test_agent_retries_with_a_rewritten_query_and_then_succeeds(self):
        bm25_only = lambda q, k: [i for i, _ in self.pipe.bm25.search(q, k)]
        ans, trace = agentic_answer("What is the cash for each group?", bm25_only, self.pipe.by_id, k=1)
        self.assertFalse(ans.abstained)
        self.assertEqual(len(trace), 2)
        self.assertIn("credit", trace[1])
        self.assertEqual(ans.citations, ["credits#0"])

    def test_rewrite_adds_shared_concept_words(self):
        self.assertIn("credit", rewrite_query("What is the cash for each group?", []).split())


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