"""Agent 池服务：注册/注销/心跳/负载/离线剔除。"""
import logging
import time

import redis.asyncio as aioredis

from app.constants import AgentStatus
from app.models.schemas import AgentHeartbeat, AgentInfo, AgentRegister
from app.repository.agent_repo import AgentRepo

logger = logging.getLogger(__name__)


class AgentService:
    def __init__(self, redis: aioredis.Redis, repo: AgentRepo | None = None):
        self.redis = redis
        self.repo = repo or AgentRepo(redis)

    async def register(self, reg: AgentRegister) -> AgentInfo:
        agent = AgentInfo(
            agent_id=reg.agent_id,
            endpoint=reg.endpoint,
            agent_tags=reg.agent_tags,
            max_concurrent=reg.max_concurrent,
            current_load=reg.current_load,
            priority=reg.priority,
            status=AgentStatus.READY,
            create_time=time.time(),
            last_heartbeat=time.time(),
        )
        created = await self.repo.upsert(agent)
        logger.info("agent registered agent_id=%s created=%s", agent.agent_id, created)
        return agent

    async def unregister(self, agent_id: str) -> bool:
        existed = await self.repo.get(agent_id)
        if not existed:
            return False
        await self.repo.remove(agent_id)
        logger.info("agent unregistered agent_id=%s", agent_id)
        return True

    async def heartbeat(self, hb: AgentHeartbeat) -> bool:
        ok = await self.repo.beat(hb.agent_id, hb.current_load)
        if not ok:
            logger.warning("heartbeat from unknown agent agent_id=%s", hb.agent_id)
        return ok

    async def online(self) -> list[AgentInfo]:
        return await self.repo.online()

    async def all(self) -> list[AgentInfo]:
        return await self.repo.all()

    async def get(self, agent_id: str) -> AgentInfo | None:
        return await self.repo.get(agent_id)

    async def set_unavailable(self, agent_id: str) -> bool:
        """后台手动置为「不可用」（unavailable），心跳不会复活。"""
        a = await self.repo.get(agent_id)
        if not a:
            return False
        await self.repo.set_unavailable(agent_id)
        logger.info("agent set unavailable agent_id=%s", agent_id)
        return True

    async def set_available(self, agent_id: str) -> bool:
        """后台手动恢复「可用」（ready），清除不可用状态。"""
        a = await self.repo.get(agent_id)
        if not a:
            return False
        await self.repo.set_available(agent_id)
        logger.info("agent set available agent_id=%s", agent_id)
        return True

    async def set_priority(self, agent_id: str, priority: int) -> AgentInfo | None:
        """设置 Agent 调度优先级（1-5，越小越先分配）。"""
        if priority < 1 or priority > 5:
            raise ValueError("priority must be in [1, 5]")
        ok = await self.repo.update_priority(agent_id, priority)
        if not ok:
            return None
        logger.info("agent priority updated agent_id=%s priority=%s", agent_id, priority)
        return await self.repo.get(agent_id)

    async def purge_stale(self, timeout: float) -> list[str]:
        """剔除超时未心跳的 Agent：非 unavailable 的置为 offline（心跳失联）。

        手动置为不可用的 agent 不参与心跳计时（已从心跳 ZSet 移除），保持 unavailable。
        """
        stale = await self.repo.stale_agents(timeout)
        purged = []
        for aid in stale:
            a = await self.repo.get(aid)
            if a and a.status == AgentStatus.UNAVAILABLE:
                continue
            await self.repo.update_status(aid, AgentStatus.OFFLINE)
            logger.info("agent purged (stale) agent_id=%s", aid)
            purged.append(aid)
        return purged