"""规则调度服务：标签匹配→负载过滤→最优选择→任务绑定→推送下发。"""
import logging

import redis.asyncio as aioredis

from app.constants import AgentStatus, TaskStatus
from app.models.schemas import AgentInfo
from app.repository.agent_repo import AgentRepo
from app.repository.task_repo import TaskRepo
from app.services.relay import RelayService

logger = logging.getLogger(__name__)


class Scheduler:
    def __init__(
        self,
        redis: aioredis.Redis,
        task_repo: TaskRepo | None = None,
        agent_repo: AgentRepo | None = None,
        relay: RelayService | None = None,
    ):
        self.redis = redis
        self.task_repo = task_repo or TaskRepo(redis)
        self.agent_repo = agent_repo or AgentRepo(redis)
        self.relay = relay or RelayService(redis)

    async def pick_candidate(self, task_tags: list[str]) -> AgentInfo | None:
        """按规则挑选最优 Agent：标签匹配→负载过滤→优先级最小且空闲。"""
        candidates = await self.agent_repo.by_tags(task_tags)
        # 负载过滤：剔除满载
        candidates = [a for a in candidates if a.current_load < a.max_concurrent]
        if not candidates:
            return None
        # 最优：优先级最小（越先分配），同优先级取负载最低、在线时间最长（稳定）
        candidates.sort(key=lambda a: (a.priority, a.current_load, -a.create_time))
        return candidates[0]

    async def dispatch(self, request_id: str) -> bool:
        """将单个 pending 任务下发给匹配 Agent（绑定 → 推送 payload → 失败回退）。"""
        task = await self.task_repo.get(request_id)
        if not task or task.status != TaskStatus.PENDING:
            return False
        agent = await self.pick_candidate(task.task_tags)
        if not agent:
            return False  # 暂无可用 Agent，保持 pending，等待下次调度
        # 绑定 Agent，标记 running
        await self.task_repo.update(
            request_id,
            status=TaskStatus.RUNNING,
            agent_id=agent.agent_id,
            progress=0,
        )
        await self.task_repo.mark_running(request_id)
        # 标记 Agent 为处理中
        await self.agent_repo.update_status(agent.agent_id, AgentStatus.PROCESSING)
        # 更新 Agent 负载
        await self.agent_repo.adjust_load(agent.agent_id, 1)
        # 真实推送 payload 到 Agent 端点
        pushed = await self.relay.dispatch_command(agent.agent_id, request_id, task.payload or {})
        if not pushed:
            # 推送失败：回退 pending 并恢复 agent 就绪、释放负载，等待下次调度（避免重复扣负载）
            await self.task_repo.update(request_id, status=TaskStatus.PENDING, agent_id="", progress=0)
            await self.task_repo.mark_pending(request_id)
            await self.agent_repo.adjust_load(agent.agent_id, -1)
            await self.agent_repo.update_status(agent.agent_id, AgentStatus.READY)
            logger.warning("task dispatch push failed, reverted request_id=%s agent=%s", request_id, agent.agent_id)
            return False
        logger.info("task dispatched request_id=%s -> agent=%s", request_id, agent.agent_id)
        return True

    async def dispatch_pending(self, batch_size: int = 20) -> int:
        """调度所有可调度的 pending 任务，返回成功下发数。"""
        pending = await self.task_repo.pending_tasks()
        dispatched = 0
        for rid in pending[:batch_size]:
            if await self.dispatch(rid):
                dispatched += 1
        return dispatched