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

与业务系统解耦：通过 session_factory 注入数据库会话，
通过 agent_model_factory 注入 ORM 模型。
"""

import logging
from datetime import datetime, timedelta, timezone
from typing import Any, Callable, Awaitable

from agentkit.core.exceptions import (
    AgentNotFoundError,
    AgentUnavailableError,
    NoAvailableAgentError,
)
from agentkit.core.protocol import AgentCapability, AgentStatus

logger = logging.getLogger(__name__)

HEARTBEAT_TIMEOUT_SECONDS = 90


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

    使用依赖注入模式，不依赖具体的 ORM 模型或数据库连接。
    """

    def __init__(
        self,
        session_factory: Callable[[], Any],
        agent_model: Any,
        load_balancer: str = "round_robin",
    ):
        """
        Args:
            session_factory: 返回 async context manager 的工厂，用于获取数据库会话
            agent_model: Agent ORM 模型类
            load_balancer: 负载均衡策略 (round_robin / least_tasks / random)
        """
        self._session_factory = session_factory
        self._agent_model = agent_model
        self._load_balancer = load_balancer
        self._round_robin_index: dict[str, int] = {}

    async def register(self, capability: AgentCapability, endpoint: str) -> str:
        """注册 Agent，返回 agent_id。同名 Agent 已存在则更新。"""
        async with self._session_factory() as db:
            try:
                Model = self._agent_model
                stmt = type(db).execute.__self__.__class__  # placeholder

                # 尝试查找已有记录
                from sqlalchemy import select
                stmt = select(Model).where(Model.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 = Model(
                        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 self._session_factory() as db:
            try:
                from sqlalchemy import select
                Model = self._agent_model
                stmt = select(Model).where(Model.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 self._session_factory() as db:
            try:
                from sqlalchemy import update
                Model = self._agent_model
                stmt = (
                    update(Model)
                    .where(Model.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 self._session_factory() as db:
            from sqlalchemy import select
            Model = self._agent_model
            stmt = select(Model).where(Model.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 self._session_factory() as db:
            from sqlalchemy import select
            Model = self._agent_model
            stmt = select(Model)
            if agent_type:
                stmt = stmt.where(Model.agent_type == agent_type)
            if status:
                stmt = stmt.where(Model.status == status)
            stmt = stmt.order_by(Model.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（支持负载均衡）"""
        async with self._session_factory() as db:
            from sqlalchemy import select
            Model = self._agent_model
            stmt = select(Model).where(Model.status == AgentStatus.ONLINE)
            result = await db.execute(stmt)
            agents = result.scalars().all()

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

            if not candidates:
                return None

            # 负载均衡选择
            if self._load_balancer == "round_robin":
                idx = self._round_robin_index.get(task_type, 0)
                selected = candidates[idx % len(candidates)]
                self._round_robin_index[task_type] = idx + 1
                return selected.name
            elif self._load_balancer == "random":
                import random
                return random.choice(candidates).name
            else:  # least_tasks 或默认：返回第一个
                return candidates[0].name

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

                stmt = (
                    update(Model)
                    .where(
                        Model.status == AgentStatus.ONLINE,
                        Model.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: Any) -> 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,
        }
