"""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,
            status=AgentStatus.ONLINE,
            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 take_offline(self, agent_id: str) -> bool:
        """手动下线 Agent。"""
        a = await self.repo.get(agent_id)
        if not a:
            return False
        await self.repo.update_status(agent_id, AgentStatus.OFFLINE)
        logger.info("agent taken offline agent_id=%s", agent_id)
        return True

    async def purge_stale(self, timeout: float) -> list[str]:
        """剔除超时未心跳的 Agent。"""
        stale = await self.repo.stale_agents(timeout)
        for aid in stale:
            await self.repo.update_status(aid, AgentStatus.OFFLINE)
            logger.info("agent purged (stale) agent_id=%s", aid)
        return stale