import unittest

from lib import cluster, schema


def make_candidate(candidate_id: str, source: str, title: str, snippet: str, score: float) -> schema.Candidate:
    return schema.Candidate(
        candidate_id=candidate_id,
        item_id=candidate_id,
        source=source,
        title=title,
        url=f"https://example.com/{candidate_id}",
        snippet=snippet,
        subquery_labels=["primary"],
        native_ranks={"primary:reddit": 1},
        local_relevance=0.8,
        freshness=80,
        engagement=10,
        source_quality=0.7,
        rrf_score=0.02,
        rerank_score=score,
        final_score=score,
    )


class ClusterV3Tests(unittest.TestCase):
    def test_singleton_clusters_for_non_clustered_plan(self):
        plan = schema.QueryPlan(
            intent="how_to",
            freshness_mode="balanced_recent",
            cluster_mode="none",
            raw_topic="docker setup",
            subqueries=[schema.SubQuery(label="primary", search_query="docker setup", ranking_query="How do I set up Docker?", sources=["reddit"])],
            source_weights={"reddit": 1.0},
        )
        candidates = [
            make_candidate("c1", "reddit", "Docker setup guide", "Step by step setup", 80),
            make_candidate("c2", "youtube", "Docker install video", "Video walkthrough", 75),
        ]
        clusters = cluster.cluster_candidates(candidates, plan)
        self.assertEqual(2, len(clusters))
        self.assertEqual(["c1"], clusters[0].representative_ids)
        self.assertEqual(["c2"], clusters[1].representative_ids)

    def test_breaking_news_clusters_related_items(self):
        plan = schema.QueryPlan(
            intent="breaking_news",
            freshness_mode="strict_recent",
            cluster_mode="story",
            raw_topic="model launch",
            subqueries=[schema.SubQuery(label="primary", search_query="model launch", ranking_query="What happened in the model launch?", sources=["reddit", "x"])],
            source_weights={"reddit": 0.5, "x": 0.5},
        )
        candidates = [
            make_candidate("c1", "reddit", "Open model launch reactions", "People are reacting to the open model launch today.", 88),
            make_candidate("c2", "x", "Open model launch update", "People are reacting to the open model launch today on X.", 84),
            make_candidate("c3", "youtube", "Different topic", "A separate discussion about hardware benchmarks.", 70),
        ]
        clusters = cluster.cluster_candidates(candidates, plan)
        self.assertEqual(2, len(clusters))
        self.assertEqual(2, len(clusters[0].candidate_ids))
        self.assertIn("c1", clusters[0].candidate_ids)
        self.assertIn("c2", clusters[0].candidate_ids)


class TestCrossSourceMerging(unittest.TestCase):
    """Test the entity-based second pass that merges same-story clusters across sources."""

    def _plan(self, intent="breaking_news"):
        return schema.QueryPlan(
            intent=intent,
            freshness_mode="strict_recent",
            cluster_mode="story",
            raw_topic="test",
            subqueries=[schema.SubQuery(label="primary", search_query="test", ranking_query="test", sources=["reddit", "x", "tiktok"])],
            source_weights={"reddit": 0.5, "x": 0.5, "tiktok": 0.5},
        )

    def test_same_story_different_phrasing_merges(self):
        """Wireless Festival example: same event, different wording, different sources."""
        candidates = [
            make_candidate("c1", "reddit", "Kanye West to headline all three nights of Wireless Festival 2026", "Big announcement for Wireless.", 80),
            make_candidate("c2", "x", "BREAKING: Kanye West is making his massive UK comeback at Wireless Festival this July", "Ye returns to UK.", 75),
            make_candidate("c3", "youtube", "Kanye West BULLY Album Review - Knox Hill Reacts", "Full album reaction and breakdown.", 70),
        ]
        clusters = cluster.cluster_candidates(candidates, self._plan())
        # c1 and c2 should merge (Kanye + Wireless + Festival overlap), c3 should stay separate
        self.assertEqual(2, len(clusters))
        wireless_cluster = next(cl for cl in clusters if len(cl.candidate_ids) == 2)
        self.assertIn("c1", wireless_cluster.candidate_ids)
        self.assertIn("c2", wireless_cluster.candidate_ids)
        self.assertEqual(sorted(["reddit", "x"]), wireless_cluster.sources)
        # Multi-source cluster should not have "single-source" uncertainty
        self.assertNotEqual("single-source", wireless_cluster.uncertainty)

    def test_different_stories_dont_merge(self):
        """Different topics should stay separate even with some entity overlap (e.g., 'Kanye')."""
        candidates = [
            make_candidate("c1", "reddit", "Kanye West BULLY Album First Impressions Thread", "What do you think of BULLY?", 80),
            make_candidate("c2", "x", "Kanye West apology for antisemitism in Wall Street Journal ad", "Full page WSJ ad.", 75),
            make_candidate("c3", "tiktok", "Kanye West Wireless Festival ticket prices breakdown", "How much for Wireless tickets?", 70),
        ]
        clusters = cluster.cluster_candidates(candidates, self._plan())
        # These are 3 different stories, should remain as 3 clusters
        self.assertEqual(3, len(clusters))

    def test_same_source_clusters_dont_merge(self):
        """Two single-source clusters from the same source should not merge via entity pass."""
        candidates = [
            make_candidate("c1", "reddit", "Kanye West Wireless Festival headline announcement", "Three nights!", 80),
            make_candidate("c2", "reddit", "Kanye West returning to Wireless Festival confirmed", "UK comeback.", 70),
        ]
        clusters = cluster.cluster_candidates(candidates, self._plan())
        # The initial greedy pass may or may not merge these (depends on token similarity).
        # But if they end up as separate clusters, the entity pass should NOT merge them
        # since they're both from reddit.
        for cl in clusters:
            self.assertTrue(len(cl.sources) >= 1)  # basic sanity


class TestPolymarketIsolation(unittest.TestCase):
    """Polymarket clusters must not merge with non-Polymarket clusters via entity overlap."""

    def _plan(self):
        return schema.QueryPlan(
            intent="breaking_news",
            freshness_mode="strict_recent",
            cluster_mode="story",
            raw_topic="test",
            subqueries=[schema.SubQuery(label="primary", search_query="test", ranking_query="test", sources=["reddit", "x", "polymarket"])],
            source_weights={"reddit": 0.5, "x": 0.5, "polymarket": 0.5},
        )

    def test_polymarket_does_not_merge_into_news_cluster(self):
        """A Polymarket prediction about Sam Altman should not merge into a news cluster about Sam Altman."""
        candidates = [
            make_candidate("c1", "reddit", "Sam Altman personal rivalry with Elon Musk escalates", "The feud between Sam Altman and Elon Musk continues.", 80),
            make_candidate("c2", "polymarket", "Sam Altman equity stake in OpenAI valued at $500M", "Will Sam Altman receive equity in OpenAI restructuring?", 75),
        ]
        clusters = cluster.cluster_candidates(candidates, self._plan())
        self.assertEqual(2, len(clusters), "Polymarket and news clusters should remain separate")
        # Each cluster should have exactly one candidate
        for cl in clusters:
            self.assertEqual(1, len(cl.candidate_ids))

    def test_two_polymarket_clusters_not_blocked_by_poly_guard(self):
        """Two Polymarket items about the same topic are not blocked by the Polymarket guard.

        Note: same-source clusters are still blocked by the existing same-source
        guard, so we verify the poly guard specifically by checking that two
        polymarket items with high text similarity merge via the greedy pass.
        """
        candidates = [
            make_candidate("c1", "polymarket", "Sam Altman equity stake in OpenAI restructuring", "Will Sam Altman get equity in the OpenAI restructuring deal?", 80),
            make_candidate("c2", "polymarket", "Sam Altman equity stake in OpenAI restructuring odds", "Will Sam Altman get equity in the OpenAI restructuring deal? Current odds.", 75),
        ]
        clusters = cluster.cluster_candidates(candidates, self._plan())
        # High text similarity means greedy pass merges them
        self.assertEqual(1, len(clusters))
        self.assertEqual(2, len(clusters[0].candidate_ids))

    def test_neither_polymarket_still_merges(self):
        """Non-Polymarket clusters with entity overlap should still merge (existing behavior)."""
        candidates = [
            make_candidate("c1", "reddit", "Sam Altman OpenAI restructuring announcement details", "Sam Altman announces major OpenAI restructuring.", 80),
            make_candidate("c2", "x", "Sam Altman reveals OpenAI restructuring plan for 2026", "Major OpenAI restructuring coming says Sam Altman.", 75),
        ]
        clusters = cluster.cluster_candidates(candidates, self._plan())
        self.assertEqual(1, len(clusters))
        self.assertEqual(2, len(clusters[0].candidate_ids))


class TestClusterUncertainty(unittest.TestCase):
    def test_single_source_returns_single_source(self):
        candidates = [make_candidate("c1", "reddit", "Title", "Body", 80)]
        result = cluster._cluster_uncertainty(candidates)
        self.assertEqual("single-source", result)

    def test_multi_source_high_score_returns_none(self):
        candidates = [
            make_candidate("c1", "reddit", "Title", "Body", 80),
            make_candidate("c2", "x", "Title2", "Body2", 70),
        ]
        result = cluster._cluster_uncertainty(candidates)
        self.assertIsNone(result)

    def test_multi_source_low_score_returns_thin_evidence(self):
        candidates = [
            make_candidate("c1", "reddit", "Title", "Body", 30),
            make_candidate("c2", "x", "Title2", "Body2", 40),
        ]
        result = cluster._cluster_uncertainty(candidates)
        self.assertEqual("thin-evidence", result)

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