from app.ai_provider.llm.pool import (
    CostAwareStrategy,
    EndpointPool,
    LeastConnectionsStrategy,
    LowestLatencyStrategy,
    PoolFactory,
    PriorityFailoverStrategy,
    WeightedRoundRobinStrategy,
)


class TestWeightedRoundRobin:
    def test_distribution(self):
        endpoints = [{"name": "a", "weight": 100}, {"name": "b", "weight": 50}]
        strategy = WeightedRoundRobinStrategy(endpoints)
        picks = [strategy.select([])["name"] for _ in range(15)]
        assert picks.count("a") > picks.count("b")

    def test_skip_excluded(self):
        endpoints = [{"name": "a", "weight": 100}, {"name": "b", "weight": 100}]
        strategy = WeightedRoundRobinStrategy(endpoints)
        pick = strategy.select(["a"])
        assert pick["name"] == "b"


class TestPriorityFailover:
    def test_primary_selected(self):
        endpoints = [{"name": "primary", "priority": 10}, {"name": "backup", "priority": 20}]
        strategy = PriorityFailoverStrategy(endpoints)
        pick = strategy.select([])
        assert pick["name"] == "primary"

    def test_failover_on_excluded(self):
        endpoints = [{"name": "primary", "priority": 10}, {"name": "backup", "priority": 20}]
        strategy = PriorityFailoverStrategy(endpoints)
        pick = strategy.select(["primary"])
        assert pick["name"] == "backup"

    def test_exhausted(self):
        endpoints = [{"name": "a", "priority": 10}]
        strategy = PriorityFailoverStrategy(endpoints)
        pick = strategy.select(["a"])
        assert pick is None


class TestLowestLatency:
    def test_pick_lowest_latency(self):
        endpoints = [{"name": "fast", "latency_ms": 100}, {"name": "slow", "latency_ms": 500}]
        strategy = LowestLatencyStrategy(endpoints)
        pick = strategy.select([])
        assert pick["name"] == "fast"


class TestCostAware:
    def test_pick_cheapest(self):
        endpoints = [{"name": "cheap"}, {"name": "expensive"}]

        def _cost(name: str) -> float:
            return 0.1 if name == "cheap" else 1.0

        strategy = CostAwareStrategy(endpoints, get_cost=_cost)
        pick = strategy.select([])
        assert pick["name"] == "cheap"


class TestLeastConnections:
    def test_pick_least_connections(self):
        endpoints = [{"name": "busy", "connections": 5}, {"name": "idle", "connections": 0}]
        strategy = LeastConnectionsStrategy(endpoints)
        pick = strategy.select([])
        assert pick["name"] == "idle"


class TestEndpointPool:
    def test_select_with_weighted_round_robin(self):
        pool = EndpointPool(
            logical_model_name="test",
            endpoints_configs=[
                {"name": "a", "weight": 100, "status": "active"},
                {"name": "b", "weight": 50, "status": "active"},
            ],
            strategy="weighted_round_robin",
        )
        selected = pool.select([])
        assert selected["name"] in ("a", "b")

    def test_select_all_excluded_returns_none(self):
        pool = EndpointPool(
            logical_model_name="test",
            endpoints_configs=[{"name": "a", "weight": 100, "status": "active"}],
        )
        assert pool.select(["a"]) is None

    def test_strategy_factory(self):
        pool = PoolFactory.create_pool(
            logical_model_name="test",
            endpoints=[{"name": "a", "weight": 100, "status": "active"}],
            strategy_name="priority_failover",
        )
        assert isinstance(pool.strategy, PriorityFailoverStrategy)