"""OrganizationContext - 组织上下文，管理 AgentProfile 与能力矩阵"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass
class AgentProfile:
    """Agent 档案 - 描述组织中一个 Agent 的能力与状态"""

    name: str
    agent_type: str  # "react", "rewoo", "plan_exec", "reflexion", "direct"
    capabilities: list[str]  # capability tag strings
    skills: list[str]  # skill names
    current_load: int = 0  # number of active tasks
    max_concurrency: int = 1
    availability: bool = True
    specializations: list[str] = field(default_factory=list)
    model: str = "default"
    execution_mode: str = "react"


class OrganizationContext:
    """组织上下文 - 管理 Agent 档案与能力矩阵，支持基于能力的 Agent 发现"""

    def __init__(self) -> None:
        self._agents: dict[str, AgentProfile] = {}
        self._capability_matrix: dict[str, list[str]] = {}  # capability -> [agent_names]

    def register_agent(self, profile: AgentProfile) -> None:
        """注册 Agent 档案"""
        self._agents[profile.name] = profile
        # 更新能力矩阵
        for cap in profile.capabilities:
            cap_lower = cap.lower()
            if cap_lower not in self._capability_matrix:
                self._capability_matrix[cap_lower] = []
            if profile.name not in self._capability_matrix[cap_lower]:
                self._capability_matrix[cap_lower].append(profile.name)
        logger.info(f"Agent profile '{profile.name}' registered")

    def unregister_agent(self, name: str) -> None:
        """注销 Agent 档案"""
        profile = self._agents.pop(name, None)
        if profile is None:
            return
        # 清理能力矩阵
        for cap in profile.capabilities:
            cap_lower = cap.lower()
            if cap_lower in self._capability_matrix:
                self._capability_matrix[cap_lower] = [
                    n for n in self._capability_matrix[cap_lower] if n != name
                ]
                if not self._capability_matrix[cap_lower]:
                    del self._capability_matrix[cap_lower]
        logger.info(f"Agent profile '{name}' unregistered")

    def get_agent_profile(self, name: str) -> AgentProfile | None:
        """获取 Agent 档案"""
        return self._agents.get(name)

    def list_agents(self) -> list[AgentProfile]:
        """列出所有 Agent 档案"""
        return list(self._agents.values())

    def find_best_agent(
        self,
        required_capabilities: list[str],
        exclude: list[str] | None = None,
    ) -> AgentProfile | None:
        """根据能力需求找到最佳 Agent

        逻辑：
        1. 找到拥有所有所需能力的 Agent
        2. 在匹配的 Agent 中，优先选择 current_load 较低的
        3. 排除 exclude 列表中的 Agent
        4. 排除不可用的 Agent
        5. 没有匹配则返回 None
        """
        exclude_set = set(exclude or [])

        # 对每个所需能力，查找拥有该能力的 Agent 名称集合
        candidate_names: set[str] | None = None
        for cap in required_capabilities:
            cap_lower = cap.lower()
            agents_with_cap = set(self._capability_matrix.get(cap_lower, []))
            if candidate_names is None:
                candidate_names = agents_with_cap
            else:
                candidate_names &= agents_with_cap

        if not candidate_names:
            return None

        # 过滤排除和不可用的 Agent，按 load 排序
        candidates = [
            self._agents[name]
            for name in candidate_names
            if name not in exclude_set
            and name in self._agents
            and self._agents[name].availability
        ]

        if not candidates:
            return None

        candidates.sort(key=lambda p: p.current_load)
        return candidates[0]

    def update_load(self, name: str, delta: int) -> None:
        """更新 Agent 负载"""
        profile = self._agents.get(name)
        if profile is not None:
            profile.current_load = max(0, profile.current_load + delta)

    def set_availability(self, name: str, available: bool) -> None:
        """设置 Agent 可用性"""
        profile = self._agents.get(name)
        if profile is not None:
            profile.availability = available

    @classmethod
    def from_agent_pool(cls, agent_pool, skill_registry) -> OrganizationContext:
        """从 AgentPool 和 SkillRegistry 构建 OrganizationContext

        Args:
            agent_pool: AgentPool 实例，提供运行时 Agent 列表
            skill_registry: SkillRegistry 实例，提供 Skill 配置查询
        """
        ctx = cls()

        if agent_pool is None or skill_registry is None:
            return ctx

        for agent_info in agent_pool.list_agents():
            agent_name = agent_info["name"]
            agent_type = agent_info.get("agent_type", "react")

            # 尝试从 skill_registry 获取 SkillConfig
            capabilities: list[str] = []
            skills: list[str] = []
            execution_mode = "react"
            model = "default"
            max_concurrency = 1

            try:
                skill = skill_registry.get(agent_name)
                config = skill.config
                capabilities = [cap.tag for cap in config.capabilities]
                execution_mode = config.execution_mode
                model = config.llm.get("model", "default") if config.llm else "default"
                max_concurrency = config.max_concurrency
                skills = [agent_name]
            except Exception:
                # Agent 不在 skill_registry 中，使用默认值
                skills = [agent_name]

            profile = AgentProfile(
                name=agent_name,
                agent_type=agent_type,
                capabilities=capabilities,
                skills=skills,
                execution_mode=execution_mode,
                model=model,
                max_concurrency=max_concurrency,
            )
            ctx.register_agent(profile)

        return ctx