from __future__ import annotations

import unittest

from agent_harness.pilot import retrieval_metrics
from agent_harness.repository import SourceFile, chunk_file
from agent_harness.retrieval import BM25FuzzyRetriever, ExactRetriever, query_terms


class RetrievalTests(unittest.TestCase):
    def setUp(self) -> None:
        sources = [
            SourceFile(
                "cache/s3.go",
                "package cache\nfunc newS3Client() { checksumValidation := whenRequired }\n",
            ),
            SourceFile(
                "commands/proxy.go",
                "package commands\nfunc executeProxy() { waitForChild() }\n",
            ),
        ]
        self.chunks = tuple(
            chunk
            for source in sources
            for chunk in chunk_file(source, chunk_lines=120, overlap_lines=20, char_limit=16000)
        )

    def test_query_terms_remove_stop_words_and_split_identifiers(self) -> None:
        terms = query_terms("Fix checksumValidation for the S3-compatible endpoint")
        self.assertIn("checksum", terms)
        self.assertIn("validation", terms)
        self.assertIn("s3", terms)
        self.assertNotIn("the", terms)

    def test_exact_and_bm25_retrieve_relevant_file(self) -> None:
        query = "S3 checksum validation fails for custom endpoints"
        exact = ExactRetriever(self.chunks).retrieve(query, 10)
        bm25 = BM25FuzzyRetriever(self.chunks).retrieve(query, 10)
        self.assertEqual(exact[0].path, "cache/s3.go")
        self.assertEqual(bm25[0].path, "cache/s3.go")

    def test_retrieval_metrics(self) -> None:
        metrics = retrieval_metrics(
            ["irrelevant.go", "gold_a.go", "gold_b.go"],
            ["gold_a.go", "gold_b.go"],
        )
        self.assertEqual(metrics["file_recall_at_1"], 0.0)
        self.assertEqual(metrics["file_recall_at_5"], 1.0)
        self.assertEqual(metrics["first_gold_rank"], 2)
        self.assertEqual(metrics["mrr"], 0.5)
        self.assertTrue(metrics["all_gold_in_top_10"])


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