"""Agent 池 Redis 存储层。"""
import time
from typing import Any

import redis.asyncio as aioredis

from app import constants as C
from app.constants import AgentStatus
from app.models.schemas import AgentInfo


class AgentRepo:
    def __init__(self, redis: aioredis.Redis):
        self.redis = redis

    @staticmethod
    def _info_key(agent_id: str) -> str:
        return C.AGENT_INFO_KEY.format(agent_id=agent_id)

    @staticmethod
    def _tag_key(tag: str) -> str:
        return C.AGENT_TAG_KEY.format(tag=tag)

    # ---------- 写入 ----------
    async def upsert(self, agent: AgentInfo, *, heartbeat: bool = False) -> bool:
        """写入 Agent 信息，返回是否新建。非心跳时为全量注册。"""
        key = self._info_key(agent.agent_id)
        existed = await self.redis.exists(key)
        if not existed:
            await self.redis.sadd(C.AGENT_ALL_SET, agent.agent_id)
            # 建立标签索引
            for tag in agent.agent_tags:
                await self.redis.sadd(self._tag_key(tag), agent.agent_id)
        mapping: dict[str, Any] = {
            "agent_id": agent.agent_id,
            "endpoint": agent.endpoint,
            "agent_tags": ",".join(agent.agent_tags),
            "max_concurrent": str(agent.max_concurrent),
            "current_load": str(agent.current_load),
            "last_heartbeat": str(agent.last_heartbeat),
            "status": agent.status.value,
            "create_time": str(agent.create_time),
        }
        await self.redis.hset(key, mapping=mapping)
        # 心跳 ZSet 记录时间戳
        await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent.agent_id: agent.last_heartbeat})
        return not existed

    async def update_status(self, agent_id: str, status: AgentStatus) -> bool:
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        await self.redis.hset(key, "status", status.value)
        if status == AgentStatus.OFFLINE:
            await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
        return True

    async def beat(self, agent_id: str, current_load: int) -> bool:
        """更新心跳时间与负载。"""
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        now = time.time()
        await self.redis.hset(
            key,
            mapping={"last_heartbeat": str(now), "current_load": str(current_load), "status": AgentStatus.ONLINE.value},
        )
        await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: now})
        return True

    async def adjust_load(self, agent_id: str, delta: int) -> bool:
        """原子调整 Agent 负载（delta 可为正/负），负载下限为 0。"""
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        # hincrby 为原子增量操作，兼容 fakeredis（不支持 Lua eval）
        new = await self.redis.hincrby(key, "current_load", delta)
        if new < 0:
            await self.redis.hset(key, "current_load", "0")
        return True

    # ---------- 读取 ----------
    async def get(self, agent_id: str) -> AgentInfo | None:
        raw = await self.redis.hgetall(self._info_key(agent_id))
        if not raw:
            return None
        return AgentInfo(
            agent_id=raw.get("agent_id", ""),
            endpoint=raw.get("endpoint", ""),
            agent_tags=[t for t in raw.get("agent_tags", "").split(",") if t],
            max_concurrent=int(raw.get("max_concurrent", 1) or 1),
            current_load=int(raw.get("current_load", 0) or 0),
            last_heartbeat=float(raw.get("last_heartbeat", 0) or 0),
            status=AgentStatus(raw.get("status", AgentStatus.ONLINE.value)),
            create_time=float(raw.get("create_time", 0) or 0),
        )

    async def online(self) -> list[AgentInfo]:
        return [a for a in await self.all() if a.status == AgentStatus.ONLINE]

    async def all(self) -> list[AgentInfo]:
        ids = list(await self.redis.smembers(C.AGENT_ALL_SET))
        out = []
        for aid in ids:
            a = await self.get(aid)
            if a:
                out.append(a)
        return out

    async def by_tags(self, tags: list[str]) -> list[AgentInfo]:
        """按标签索引取 Agent 池（取所有标签的交集，无标签则返回全部在线）。"""
        if not tags:
            return await self.online()
        keys = [self._tag_key(t) for t in tags]
        if len(keys) == 1:
            ids = list(await self.redis.smembers(keys[0]))
        else:
            await self.redis.sinterstore("agent:tmp:intersect", keys)
            ids = list(await self.redis.smembers("agent:tmp:intersect"))
            await self.redis.delete("agent:tmp:intersect")
        out = []
        for aid in ids:
            a = await self.get(aid)
            if a and a.status == AgentStatus.ONLINE:
                out.append(a)
        return out

    async def remove(self, agent_id: str) -> None:
        """注销：删除信息、心跳 ZSet、标签索引、全量索引。"""
        a = await self.get(agent_id)
        if a:
            for tag in a.agent_tags:
                await self.redis.srem(self._tag_key(tag), agent_id)
        await self.redis.delete(self._info_key(agent_id))
        await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
        await self.redis.srem(C.AGENT_ALL_SET, agent_id)

    async def stale_agents(self, timeout: float) -> list[str]:
        """返回超过 timeout 秒未心跳的 Agent ID（基于 ZSet score）。"""
        cutoff = time.time() - timeout
        scored = await self.redis.zrangebyscore(C.AGENT_HEARTBEAT_ZSET, 0, cutoff)
        return list(scored)

    async def heartbeat_map(self) -> dict[str, float]:
        """返回 {agent_id: last_heartbeat}。"""
        return {k: float(v) for k, v in await self.redis.zrange(C.AGENT_HEARTBEAT_ZSET, 0, -1, withscores=True)}