"""Agent 注册中心 - 管理 Agent 的注册、发现、状态"""

import logging
import uuid
from datetime import datetime, timedelta, timezone

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.agent_framework.exceptions import (
    AgentAlreadyRegisteredError,
    AgentNotFoundError,
    AgentUnavailableError,
    NoAvailableAgentError,
)
from app.agent_framework.protocol import AgentCapability, AgentStatus
from app.database import AsyncSessionLocal
from app.models.agent import AgentRegistry as AgentRegistryModel

logger = logging.getLogger(__name__)

# 心跳超时阈值（秒）
HEARTBEAT_TIMEOUT_SECONDS = 90


class AgentRegistry:
    """Agent 注册中心，管理 Agent 的注册、发现、状态"""

    async def register(self, capability: AgentCapability, endpoint: str) -> str:
        """
        注册 Agent，返回 agent_id。
        如果同名 Agent 已存在则更新，否则新建。
        """
        async with AsyncSessionLocal() as db:
            try:
                # 检查是否已存在
                stmt = select(AgentRegistryModel).where(
                    AgentRegistryModel.name == capability.agent_name
                )
                result = await db.execute(stmt)
                existing = result.scalar_one_or_none()

                if existing:
                    # 更新已有记录
                    existing.agent_type = capability.agent_type
                    existing.version = capability.version
                    existing.endpoint = endpoint
                    existing.description = capability.description
                    existing.capabilities = capability.to_dict()
                    existing.status = AgentStatus.ONLINE
                    existing.last_heartbeat = datetime.now(timezone.utc)
                    await db.commit()
                    await db.refresh(existing)
                    agent_id = existing.id
                    logger.info(f"Agent '{capability.agent_name}' re-registered (id={agent_id})")
                else:
                    # 新建记录
                    agent = AgentRegistryModel(
                        name=capability.agent_name,
                        display_name=capability.agent_name.replace("_", " ").title(),
                        agent_type=capability.agent_type,
                        description=capability.description,
                        version=capability.version,
                        endpoint=endpoint,
                        status=AgentStatus.ONLINE,
                        capabilities=capability.to_dict(),
                        last_heartbeat=datetime.now(timezone.utc),
                    )
                    db.add(agent)
                    await db.commit()
                    await db.refresh(agent)
                    agent_id = agent.id
                    logger.info(f"Agent '{capability.agent_name}' registered (id={agent_id})")

                return str(agent_id)

            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to register agent '{capability.agent_name}': {e}")
                raise

    async def unregister(self, agent_name: str):
        """注销 Agent（设置状态为 offline）"""
        async with AsyncSessionLocal() as db:
            try:
                stmt = select(AgentRegistryModel).where(
                    AgentRegistryModel.name == agent_name
                )
                result = await db.execute(stmt)
                agent = result.scalar_one_or_none()

                if not agent:
                    logger.warning(f"Attempted to unregister non-existent agent '{agent_name}'")
                    return

                agent.status = AgentStatus.OFFLINE
                await db.commit()
                logger.info(f"Agent '{agent_name}' unregistered")

            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to unregister agent '{agent_name}': {e}")
                raise

    async def update_heartbeat(self, agent_name: str):
        """更新心跳时间"""
        async with AsyncSessionLocal() as db:
            try:
                stmt = (
                    update(AgentRegistryModel)
                    .where(AgentRegistryModel.name == agent_name)
                    .values(
                        last_heartbeat=datetime.now(timezone.utc),
                        status=AgentStatus.ONLINE,
                    )
                )
                await db.execute(stmt)
                await db.commit()
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to update heartbeat for agent '{agent_name}': {e}")

    async def get_agent(self, agent_name: str) -> dict | None:
        """获取 Agent 信息"""
        async with AsyncSessionLocal() as db:
            stmt = select(AgentRegistryModel).where(
                AgentRegistryModel.name == agent_name
            )
            result = await db.execute(stmt)
            agent = result.scalar_one_or_none()

            if not agent:
                return None

            return self._agent_to_dict(agent)

    async def list_agents(
        self,
        agent_type: str | None = None,
        status: str | None = None,
    ) -> list[dict]:
        """列出 Agent，支持按类型和状态筛选"""
        async with AsyncSessionLocal() as db:
            stmt = select(AgentRegistryModel)
            if agent_type:
                stmt = stmt.where(AgentRegistryModel.agent_type == agent_type)
            if status:
                stmt = stmt.where(AgentRegistryModel.status == status)
            stmt = stmt.order_by(AgentRegistryModel.created_at.desc())

            result = await db.execute(stmt)
            agents = result.scalars().all()

            return [self._agent_to_dict(a) for a in agents]

    async def get_available_agent(self, task_type: str) -> str | None:
        """根据任务类型找到可用 Agent（返回 agent_name）"""
        async with AsyncSessionLocal() as db:
            # 查找状态为 online 的 Agent
            stmt = select(AgentRegistryModel).where(
                AgentRegistryModel.status == AgentStatus.ONLINE,
            )
            result = await db.execute(stmt)
            agents = result.scalars().all()

            for agent in agents:
                capabilities = agent.capabilities or {}
                supported_tasks = capabilities.get("supported_tasks", [])
                if task_type in supported_tasks:
                    return agent.name

            return None

    async def check_health(self):
        """检查所有 Agent 健康状态，超时标记为 offline"""
        async with AsyncSessionLocal() as db:
            try:
                timeout_threshold = datetime.now(timezone.utc) - timedelta(
                    seconds=HEARTBEAT_TIMEOUT_SECONDS
                )

                # 将心跳超时的 Agent 标记为 offline
                stmt = (
                    update(AgentRegistryModel)
                    .where(
                        AgentRegistryModel.status == AgentStatus.ONLINE,
                        AgentRegistryModel.last_heartbeat < timeout_threshold,
                    )
                    .values(status=AgentStatus.OFFLINE)
                )
                result = await db.execute(stmt)
                await db.commit()

                if result.rowcount > 0:
                    logger.warning(
                        f"Marked {result.rowcount} agent(s) as offline due to heartbeat timeout"
                    )

            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to check agent health: {e}")

    def _agent_to_dict(self, agent: AgentRegistryModel) -> dict:
        """将 Agent ORM 对象转换为字典"""
        return {
            "id": str(agent.id),
            "name": agent.name,
            "display_name": agent.display_name,
            "agent_type": agent.agent_type,
            "description": agent.description,
            "version": agent.version,
            "endpoint": agent.endpoint,
            "status": agent.status,
            "capabilities": agent.capabilities,
            "last_heartbeat": agent.last_heartbeat.isoformat() if agent.last_heartbeat else None,
            "created_at": agent.created_at.isoformat() if agent.created_at else None,
            "updated_at": agent.updated_at.isoformat() if agent.updated_at else None,
        }
