"""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),
            "priority": str(agent.priority),
            "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)
        # 离线（心跳失联）与不可用均退出心跳 ZSet，避免被误判/干扰
        if status in (AgentStatus.OFFLINE, AgentStatus.UNAVAILABLE):
            await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
        return True

    async def set_unavailable(self, agent_id: str) -> bool:
        """后台手动置为不可用（unavailable），心跳不会复活。"""
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        await self.redis.hset(key, "status", AgentStatus.UNAVAILABLE.value)
        await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
        return True

    async def set_available(self, agent_id: str) -> bool:
        """后台手动恢复可用（ready），清除不可用状态。"""
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        await self.redis.hset(key, "status", AgentStatus.READY.value)
        # 恢复就绪时补齐心跳时间戳，避免被误判为超时
        await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: time.time()})
        return True

    async def mark_idle(self, agent_id: str) -> bool:
        """任务结束/取消后标记 Agent 空闲。

        若 agent 已被手动置为不可用（unavailable），保持不可用、不复活；
        否则置回 ready。
        """
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        cur = await self.get(agent_id)
        if cur is None:
            return False
        if cur.status == AgentStatus.UNAVAILABLE:
            return True
        await self.redis.hset(key, "status", AgentStatus.READY.value)
        return True

    async def beat(self, agent_id: str, current_load: int) -> bool:
        """更新心跳时间与负载，按状态机规则流转状态。

        - 不可用（unavailable）：保持不可用，不因心跳复活。
        - 离线（offline，心跳失联）：恢复 ready。
        - ready/processing/stopping：保持原状态。
        """
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        now = time.time()
        cur = await self.get(agent_id)
        if cur is None:
            return False
        if cur.status == AgentStatus.UNAVAILABLE:
            status = AgentStatus.UNAVAILABLE
        elif cur.status == AgentStatus.OFFLINE:
            status = AgentStatus.READY
        else:
            status = cur.status
        await self.redis.hset(
            key,
            mapping={
                "last_heartbeat": str(now),
                "current_load": str(current_load),
                "status": status.value,
            },
        )
        await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: now})
        return True

    async def update_priority(self, agent_id: str, priority: int) -> bool:
        """更新 Agent 调度优先级。"""
        key = self._info_key(agent_id)
        if not await self.redis.exists(key):
            return False
        await self.redis.hset(key, "priority", str(priority))
        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),
            priority=int(raw.get("priority", 3) or 3),
            last_heartbeat=float(raw.get("last_heartbeat", 0) or 0),
            status=AgentStatus(raw.get("status", AgentStatus.READY.value)),
            create_time=float(raw.get("create_time", 0) or 0),
        )

    async def online(self) -> list[AgentInfo]:
        # 可参与调度的状态：ready（空闲）与 processing（忙碌但仍有并发容量）。
        # offline（心跳失联）、unavailable（手动禁用）与 stopping（正在停止）不参与调度。
        return [a for a in await self.all() if a.status in (AgentStatus.READY, AgentStatus.PROCESSING)]

    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 池（取所有标签的交集，无标签则返回全部可调度）。

        只返回可调度的 agent（ready/processing），offline/unavailable/stopping 被排除。
        """
        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 in (AgentStatus.READY, AgentStatus.PROCESSING):
                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)}