import unittest

from swarm_minimal.core import ConsensusAgent, ConsensusSwarm, ConsensusVote


def vote(agent_id: str, role: str, candidate: str, confidence: float, evidence: str) -> ConsensusVote:
    return ConsensusVote(
        agent_id=agent_id,
        role=role,
        candidate=candidate,
        confidence=confidence,
        evidence=evidence,
    )


class ConsensusConvergenceTest(unittest.TestCase):
    def test_multi_round_convergence_requires_threshold_and_margin(self) -> None:
        agents = [
            ConsensusAgent(
                id="architecture-agnet",
                role="architecture",
                weight=1.0,
                vote=lambda scores, state, round_index: vote(
                    "architecture-agnet",
                    "architecture",
                    "lease_based_pg_queue",
                    0.44 if round_index == 1 else 0.67,
                    "prefers explicit leases in PostgreSQL task pool",
                ),
            ),
            ConsensusAgent(
                id="reliability-agnet",
                role="reliability",
                weight=1.2,
                vote=lambda scores, state, round_index: vote(
                    "reliability-agnet",
                    "reliability",
                    "lease_based_pg_queue",
                    0.38 if round_index == 1 else 0.72,
                    "observed retry and lease recovery advantages",
                ),
            ),
            ConsensusAgent(
                id="latency-agnet",
                role="latency",
                weight=0.8,
                vote=lambda scores, state, round_index: vote(
                    "latency-agnet",
                    "latency",
                    "redis_only_queue" if round_index == 1 else state.get("active_candidate", "lease_based_pg_queue"),
                    0.60 if round_index == 1 else 0.58,
                    "starts from latency, then follows shared evidence",
                ),
            ),
        ]

        result = ConsensusSwarm(
            agents,
            threshold=0.70,
            min_margin=0.25,
            max_rounds=4,
            evaporation=0.9,
        ).run("choose the minimal durable scheduling strategy")

        self.assertTrue(result.converged)
        self.assertEqual(result.accepted_candidate, "lease_based_pg_queue")
        self.assertGreaterEqual(len(result.rounds), 2)
        self.assertFalse(result.rounds[0].converged)
        self.assertTrue(result.rounds[-1].converged)
        self.assertGreaterEqual(result.rounds[-1].leader_share, 0.70)
        self.assertGreaterEqual(result.rounds[-1].margin, 0.25)

    def test_distinct_agents_are_defined_by_role_weight_and_policy_not_just_count(self) -> None:
        agents = [
            ConsensusAgent(
                id="planner",
                role="plan-quality",
                weight=1.0,
                vote=lambda scores, state, round_index: vote("planner", "plan-quality", "candidate-a", 0.8, "plan coverage"),
            ),
            ConsensusAgent(
                id="verifier",
                role="verification",
                weight=1.4,
                vote=lambda scores, state, round_index: vote("verifier", "verification", "candidate-b", 0.7, "test evidence"),
            ),
        ]

        self.assertEqual({agent.role for agent in agents}, {"plan-quality", "verification"})
        self.assertNotEqual(agents[0].weight, agents[1].weight)
        result = ConsensusSwarm(agents, threshold=0.5, min_margin=0.01, max_rounds=1).run("role diversity")
        self.assertEqual(result.accepted_candidate, "candidate-b")


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