"""
Agent Capability Registry - Track and query agent capabilities.

Provides:
- Capability enumeration
- Agent capability registration
- Capability-based agent lookup
- Capability scoring for task assignment
"""

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional

logger = logging.getLogger(__name__)


class AgentCapability(Enum):
    """Standard agent capabilities."""

    # File operations
    FILE_READ = "file_read"
    FILE_WRITE = "file_write"
    FILE_DELETE = "file_delete"

    # Code operations
    CODE_EDIT = "code_edit"
    CODE_REVIEW = "code_review"
    CODE_REFACTOR = "code_refactor"
    CODE_GENERATE = "code_generate"

    # Testing
    RUN_TESTS = "run_tests"
    DEBUG = "debug"

    # Version control
    GIT = "git"
    GIT_COMMIT = "git_commit"
    GIT_MERGE = "git_merge"

    # Terminal
    TERMINAL = "terminal"
    SHELL_COMMAND = "shell_command"

    # Deployment
    DEPLOY = "deploy"
    DOCKER = "docker"
    KUBERNETES = "kubernetes"

    # Search/Research
    SEARCH = "search"
    WEB_BROWSE = "web_browse"

    # Communication
    STREAMING = "streaming"
    FUNCTION_CALLING = "function_calling"

    # Performance
    FAST = "fast"
    LARGE_CONTEXT = "large_context"
    AUTONOMOUS = "autonomous"


@dataclass
class CapabilityScore:
    """Score for a specific capability."""

    capability: str
    score: float  # 0.0 to 1.0
    notes: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "capability": self.capability,
            "score": self.score,
            "notes": self.notes,
        }


@dataclass
class AgentCapabilities:
    """Capabilities for a specific agent."""

    agent_id: str
    capabilities: set[str] = field(default_factory=set)
    scores: dict[str, CapabilityScore] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)

    def has_capability(self, capability: str | AgentCapability) -> bool:
        """Check if agent has a capability."""
        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        return cap_str in self.capabilities

    def get_score(self, capability: str | AgentCapability) -> float:
        """Get the score for a capability (default 1.0 if present, 0.0 if not)."""
        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        if cap_str in self.scores:
            return self.scores[cap_str].score
        return 1.0 if cap_str in self.capabilities else 0.0

    def add_capability(
        self,
        capability: str | AgentCapability,
        score: float = 1.0,
        notes: str = "",
    ) -> None:
        """Add a capability to this agent."""
        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        self.capabilities.add(cap_str)
        self.scores[cap_str] = CapabilityScore(cap_str, score, notes)

    def remove_capability(self, capability: str | AgentCapability) -> None:
        """Remove a capability from this agent."""
        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        self.capabilities.discard(cap_str)
        if cap_str in self.scores:
            del self.scores[cap_str]

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "agent_id": self.agent_id,
            "capabilities": list(self.capabilities),
            "scores": {k: v.to_dict() for k, v in self.scores.items()},
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "AgentCapabilities":
        """Create from dictionary."""
        caps = cls(
            agent_id=data["agent_id"],
            capabilities=set(data.get("capabilities", [])),
            metadata=data.get("metadata", {}),
        )
        for cap_data in data.get("scores", {}).values():
            caps.scores[cap_data["capability"]] = CapabilityScore(
                capability=cap_data["capability"],
                score=cap_data["score"],
                notes=cap_data.get("notes", ""),
            )
        return caps


class CapabilityRegistry:
    """
    Registry for tracking agent capabilities.

    Provides capability-based agent lookup and scoring for
    intelligent task assignment.

    Example:
        registry = CapabilityRegistry()

        # Register agent capabilities
        registry.register(
            "claude-code",
            [AgentCapability.CODE_EDIT, AgentCapability.GIT, AgentCapability.TERMINAL],
        )

        # Find agents with a capability
        agents = registry.get_agents_with_capability(AgentCapability.CODE_EDIT)

        # Find best agent for multiple capabilities
        best = registry.find_best_agent([AgentCapability.CODE_EDIT, AgentCapability.RUN_TESTS])
    """

    def __init__(self) -> None:
        """Initialize the capability registry."""
        self._agents: dict[str, AgentCapabilities] = {}
        self._capability_index: dict[str, set[str]] = {}  # capability -> set of agent_ids

    def register(
        self,
        agent_id: str,
        capabilities: list[str | AgentCapability],
        scores: Optional[dict[str, float]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> AgentCapabilities:
        """
        Register an agent with its capabilities.

        Args:
            agent_id: Unique identifier for the agent
            capabilities: List of capabilities
            scores: Optional scores for each capability (0.0-1.0)
            metadata: Optional metadata

        Returns:
            The AgentCapabilities object
        """
        scores = scores or {}
        metadata = metadata or {}

        agent_caps = AgentCapabilities(
            agent_id=agent_id,
            metadata=metadata,
        )

        for cap in capabilities:
            cap_str = cap.value if isinstance(cap, AgentCapability) else cap
            score = scores.get(cap_str, 1.0)
            agent_caps.add_capability(cap_str, score)

            # Update index
            if cap_str not in self._capability_index:
                self._capability_index[cap_str] = set()
            self._capability_index[cap_str].add(agent_id)

        self._agents[agent_id] = agent_caps
        logger.debug(f"Registered capabilities for {agent_id}: {list(agent_caps.capabilities)}")

        return agent_caps

    def unregister(self, agent_id: str) -> bool:
        """
        Unregister an agent.

        Returns True if agent was registered.
        """
        if agent_id not in self._agents:
            return False

        agent_caps = self._agents[agent_id]

        # Update index
        for cap in agent_caps.capabilities:
            if cap in self._capability_index:
                self._capability_index[cap].discard(agent_id)

        del self._agents[agent_id]
        return True

    def get_capabilities(self, agent_id: str) -> Optional[AgentCapabilities]:
        """Get capabilities for an agent."""
        return self._agents.get(agent_id)

    def has_capability(self, agent_id: str, capability: str | AgentCapability) -> bool:
        """Check if an agent has a specific capability."""
        agent_caps = self._agents.get(agent_id)
        if not agent_caps:
            return False
        return agent_caps.has_capability(capability)

    def get_agents_with_capability(
        self, capability: str | AgentCapability
    ) -> list[str]:
        """
        Get all agents that have a specific capability.

        Args:
            capability: The capability to search for

        Returns:
            List of agent IDs
        """
        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        return list(self._capability_index.get(cap_str, set()))

    def get_agents_with_all_capabilities(
        self, capabilities: list[str | AgentCapability]
    ) -> list[str]:
        """
        Get agents that have ALL specified capabilities.

        Args:
            capabilities: List of required capabilities

        Returns:
            List of agent IDs
        """
        if not capabilities:
            return list(self._agents.keys())

        cap_strs = [
            c.value if isinstance(c, AgentCapability) else c
            for c in capabilities
        ]

        result = []
        for agent_id, agent_caps in self._agents.items():
            if all(agent_caps.has_capability(c) for c in cap_strs):
                result.append(agent_id)

        return result

    def get_agents_with_any_capability(
        self, capabilities: list[str | AgentCapability]
    ) -> list[str]:
        """
        Get agents that have ANY of the specified capabilities.

        Args:
            capabilities: List of capabilities

        Returns:
            List of agent IDs
        """
        if not capabilities:
            return []

        cap_strs = [
            c.value if isinstance(c, AgentCapability) else c
            for c in capabilities
        ]

        result = set()
        for cap in cap_strs:
            result.update(self._capability_index.get(cap, set()))

        return list(result)

    def find_best_agent(
        self,
        required_capabilities: list[str | AgentCapability],
        preferred_capabilities: Optional[list[str | AgentCapability]] = None,
        excluded_agents: Optional[list[str]] = None,
    ) -> Optional[str]:
        """
        Find the best agent for a set of capabilities.

        Scoring:
        1. Must have all required capabilities
        2. Score based on capability scores and preferred capabilities

        Args:
            required_capabilities: Capabilities the agent must have
            preferred_capabilities: Nice-to-have capabilities
            excluded_agents: Agent IDs to exclude

        Returns:
            Best agent ID or None if no match
        """
        excluded = set(excluded_agents or [])
        preferred = preferred_capabilities or []

        # Get candidates with all required capabilities
        candidates = self.get_agents_with_all_capabilities(required_capabilities)
        candidates = [c for c in candidates if c not in excluded]

        if not candidates:
            return None

        # Score candidates
        scores: dict[str, float] = {}
        for agent_id in candidates:
            agent_caps = self._agents[agent_id]
            score = 0.0

            # Score from required capabilities
            for cap in required_capabilities:
                score += agent_caps.get_score(cap)

            # Bonus for preferred capabilities
            for cap in preferred:
                if agent_caps.has_capability(cap):
                    score += agent_caps.get_score(cap) * 0.5

            scores[agent_id] = score

        # Return highest scoring agent
        return max(candidates, key=lambda a: scores[a])

    def rank_agents(
        self,
        capabilities: list[str | AgentCapability],
        limit: int = 5,
    ) -> list[tuple[str, float]]:
        """
        Rank agents by their capability scores.

        Args:
            capabilities: Capabilities to score on
            limit: Maximum number of results

        Returns:
            List of (agent_id, score) tuples, sorted by score descending
        """
        cap_strs = [
            c.value if isinstance(c, AgentCapability) else c
            for c in capabilities
        ]

        scores: list[tuple[str, float]] = []
        for agent_id, agent_caps in self._agents.items():
            score = sum(agent_caps.get_score(c) for c in cap_strs)
            if score > 0:
                scores.append((agent_id, score))

        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:limit]

    def list_all_capabilities(self) -> list[str]:
        """Get all unique capabilities across all agents."""
        return list(self._capability_index.keys())

    def list_agents(self) -> list[str]:
        """Get all registered agent IDs."""
        return list(self._agents.keys())

    def get_summary(self) -> dict[str, Any]:
        """Get a summary of the registry."""
        return {
            "total_agents": len(self._agents),
            "total_capabilities": len(self._capability_index),
            "agents": {
                agent_id: list(caps.capabilities)
                for agent_id, caps in self._agents.items()
            },
            "capability_coverage": {
                cap: len(agents)
                for cap, agents in self._capability_index.items()
            },
        }

    def add_capability_to_agent(
        self,
        agent_id: str,
        capability: str | AgentCapability,
        score: float = 1.0,
    ) -> bool:
        """
        Add a capability to an existing agent.

        Returns True if added successfully.
        """
        if agent_id not in self._agents:
            return False

        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        self._agents[agent_id].add_capability(cap_str, score)

        if cap_str not in self._capability_index:
            self._capability_index[cap_str] = set()
        self._capability_index[cap_str].add(agent_id)

        return True

    def remove_capability_from_agent(
        self,
        agent_id: str,
        capability: str | AgentCapability,
    ) -> bool:
        """
        Remove a capability from an agent.

        Returns True if removed successfully.
        """
        if agent_id not in self._agents:
            return False

        cap_str = capability.value if isinstance(capability, AgentCapability) else capability
        self._agents[agent_id].remove_capability(cap_str)

        if cap_str in self._capability_index:
            self._capability_index[cap_str].discard(agent_id)

        return True
