"""规则调度服务：标签匹配→负载过滤→最优选择→任务绑定→推送下发。"""
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 → 失败回退）。

        成功 / 无可用 Agent / 推送失败时，都会在任务上记录 error_info（成功时清空），
        便于前端与日志定位“为什么没分配”。
        """
        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:
            # 暂无可用 Agent（标签不匹配 / 离线 / 满载），保持 pending，等待下次调度
            await self.task_repo.update(
                request_id,
                error_info="无可用 Agent：标签不匹配或全部离线/满载",
            )
            logger.info("task no candidate request_id=%s tags=%s", request_id, task.task_tags)
            return False
        # 绑定 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)
        # 记录调度前状态，供推送失败回退时恢复
        prev_status = agent.status
        # 标记 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, reason = 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,
                error_info=f"推送失败: {reason}",
            )
            await self.task_repo.mark_pending(request_id)
            await self.agent_repo.adjust_load(agent.agent_id, -1)
            # 若该 Agent 本为 PROCESSING（还有并发容量），保持 PROCESSING；否则置回 READY
            restore = prev_status if prev_status == AgentStatus.PROCESSING else AgentStatus.READY
            await self.agent_repo.update_status(agent.agent_id, restore)
            logger.warning("task dispatch push failed, reverted request_id=%s agent=%s reason=%s",
                           request_id, agent.agent_id, reason)
            return False
        # 推送成功：清空 error_info
        await self.task_repo.update(request_id, error_info="")
        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