from __future__ import annotations

import random
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any, Protocol


@dataclass
class PoolContext:
    request_id: str = ""
    excluded_endpoints: set[str] = field(default_factory=set)


class SchedulingStrategy(Protocol):
    def select(self, excluded: set[str]) -> dict[str, Any] | None: ...


class WeightedRoundRobinStrategy:
    """策略 1：在所有活跃端点间加权轮询。"""

    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
        self._endpoints = endpoints
        self._index = 0

    def select(self, excluded: set[str]) -> dict[str, Any] | None:
        available = [
            e
            for e in self._endpoints
            if e.get("name") not in excluded and e.get("status", "active") == "active"
        ]
        if not available:
            return None
        weights = [e.get("weight", 100) for e in available]
        self._index = (self._index + 1) % len(available)
        return random.choices(available, weights=weights, k=1)[0]


class PriorityFailoverStrategy:
    """策略 2：选择最高优先级；当排除时顺延。"""

    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
        self._endpoints = sorted(endpoints, key=lambda e: e.get("priority", 100))

    def select(self, excluded: set[str]) -> dict[str, Any] | None:
        for ep in self._endpoints:
            if ep.get("name") not in excluded and ep.get("status", "active") == "active":
                return ep
        return None


class LowestLatencyStrategy:
    """策略 3：选择最近延迟最低的端点。"""

    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
        self._endpoints = endpoints

    def select(self, excluded: set[str]) -> dict[str, Any] | None:
        available = [
            e
            for e in self._endpoints
            if e.get("name") not in excluded and e.get("status", "active") == "active"
        ]
        if not available:
            return None
        return min(available, key=lambda e: e.get("latency_ms", 1000))


class LeastConnectionsStrategy:
    """策略 4：选择活跃连接数最少的端点。"""

    def __init__(self, endpoints: list[dict[str, Any]]) -> None:
        self._endpoints = endpoints

    def select(self, excluded: set[str]) -> dict[str, Any] | None:
        available = [
            e
            for e in self._endpoints
            if e.get("name") not in excluded and e.get("status", "active") == "active"
        ]
        if not available:
            return None
        return min(available, key=lambda e: e.get("connections", 0))


class CostAwareStrategy:
    """策略 5：按每 token 成本选择最便宜的端点。"""

    def __init__(
        self, endpoints: list[dict[str, Any]], get_cost: Callable[[str], float] | None = None
    ) -> None:
        self._endpoints = endpoints
        self._get_cost = get_cost or (lambda _: 0.0)

    def select(self, excluded: set[str]) -> dict[str, Any] | None:
        available = [
            e
            for e in self._endpoints
            if e.get("name") not in excluded and e.get("status", "active") == "active"
        ]
        if not available:
            return None
        return min(available, key=lambda e: self._get_cost(e.get("name", "")))


_STRATEGIES = {
    "weighted_round_robin": WeightedRoundRobinStrategy,
    "priority_failover": PriorityFailoverStrategy,
    "lowest_latency": LowestLatencyStrategy,
    "least_connections": LeastConnectionsStrategy,
    "cost_aware": CostAwareStrategy,
}


class EndpointPool:
    """模型池，带可插拔的调度策略。"""

    def __init__(
        self,
        logical_model_name: str,
        endpoints_configs: list[dict[str, Any]],
        strategy: str = "weighted_round_robin",
    ) -> None:
        self.logical_model_name = logical_model_name
        self._endpoints = endpoints_configs
        strategy_cls = _STRATEGIES.get(strategy, WeightedRoundRobinStrategy)
        self.strategy = (
            strategy_cls(endpoints_configs)
            if strategy == "priority_failover"
            else strategy_cls(endpoints_configs)
        )

    def select(self, excluded: set[str]) -> dict[str, Any] | None:
        """使用配置的策略选择一个端点。"""
        return self.strategy.select(excluded)

    def list_endpoints(self) -> list[dict[str, Any]]:
        return list(self._endpoints)


class PoolFactory:
    """工厂：创建端点池。"""

    @staticmethod
    def create_pool(
        logical_model_name: str,
        endpoints: list[dict[str, Any]],
        strategy_name: str = "weighted_round_robin",
    ) -> EndpointPool:
        return EndpointPool(
            logical_model_name=logical_model_name,
            endpoints_configs=endpoints,
            strategy=strategy_name,
        )

    @staticmethod
    def create_strategy(
        name: str,
        endpoints: list[dict[str, Any]],
    ) -> SchedulingStrategy:
        strategy_cls = _STRATEGIES.get(name)
        if strategy_cls is None:
            raise ValueError(f"Unknown scheduling strategy: {name}")
        return strategy_cls(endpoints)