# -*- coding: utf-8 -*-
"""Agent Registry - In-memory registry for agent management."""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class AgentInfo:
    """Information about a registered agent."""

    agent_id: str
    agent_type: str
    workspace_id: str
    config_path: str
    agent_dir: str
    status: str = "inactive"  # inactive, active, error
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "agent_id": self.agent_id,
            "agent_type": self.agent_type,
            "workspace_id": self.workspace_id,
            "config_path": self.config_path,
            "agent_dir": self.agent_dir,
            "status": self.status,
            "metadata": self.metadata,
        }


class AgentRegistry:
    """In-memory registry for agent instances."""

    def __init__(self):
        """Initialize the agent registry."""
        # Dictionary mapping agent_id -> AgentInfo
        self._agents: Dict[str, AgentInfo] = {}
        # Index mapping workspace_id -> set of agent_ids
        self._workspace_index: Dict[str, set] = {}

    def register(
        self,
        agent_id: str,
        agent_type: str,
        workspace_id: str,
        config_path: str,
        agent_dir: str,
        status: str = "inactive",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> AgentInfo:
        """Register an agent in the registry.

        Args:
            agent_id: Unique identifier for the agent
            agent_type: Type of agent
            workspace_id: ID of the workspace containing the agent
            config_path: Path to agent configuration file
            agent_dir: Path to agent directory
            status: Initial status (default: inactive)
            metadata: Optional metadata dictionary

        Returns:
            AgentInfo instance

        Raises:
            ValueError: If agent_id is already registered
        """
        if agent_id in self._agents:
            raise ValueError(f"Agent '{agent_id}' is already registered")

        agent_info = AgentInfo(
            agent_id=agent_id,
            agent_type=agent_type,
            workspace_id=workspace_id,
            config_path=config_path,
            agent_dir=agent_dir,
            status=status,
            metadata=metadata or {},
        )

        self._agents[agent_id] = agent_info

        # Update workspace index
        if workspace_id not in self._workspace_index:
            self._workspace_index[workspace_id] = set()
        self._workspace_index[workspace_id].add(agent_id)

        return agent_info

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

        Args:
            agent_id: ID of the agent to unregister

        Returns:
            True if unregistered, False if agent wasn't registered
        """
        if agent_id not in self._agents:
            return False

        agent_info = self._agents[agent_id]

        # Remove from workspace index
        workspace_id = agent_info.workspace_id
        if workspace_id in self._workspace_index:
            self._workspace_index[workspace_id].discard(agent_id)
            if not self._workspace_index[workspace_id]:
                del self._workspace_index[workspace_id]

        # Remove from agents dict
        del self._agents[agent_id]

        return True

    def get(self, agent_id: str) -> Optional[AgentInfo]:
        """Get agent information by ID.

        Args:
            agent_id: ID of the agent

        Returns:
            AgentInfo if found, None otherwise
        """
        return self._agents.get(agent_id)

    def list_all(
        self,
        workspace_id: Optional[str] = None,
        agent_type: Optional[str] = None,
        status: Optional[str] = None,
    ) -> List[AgentInfo]:
        """List all registered agents with optional filtering.

        Args:
            workspace_id: Filter by workspace ID
            agent_type: Filter by agent type
            status: Filter by status

        Returns:
            List of AgentInfo instances
        """
        agents = list(self._agents.values())

        if workspace_id:
            agent_ids = self._workspace_index.get(workspace_id, set())
            agents = [a for a in agents if a.agent_id in agent_ids]

        if agent_type:
            agents = [a for a in agents if a.agent_type == agent_type]

        if status:
            agents = [a for a in agents if a.status == status]

        return agents

    def update_status(self, agent_id: str, status: str) -> bool:
        """Update the status of an agent.

        Args:
            agent_id: ID of the agent
            status: New status value

        Returns:
            True if updated, False if agent not found
        """
        if agent_id not in self._agents:
            return False

        self._agents[agent_id].status = status
        return True

    def update_metadata(self, agent_id: str, metadata: Dict[str, Any]) -> bool:
        """Update the metadata of an agent.

        Args:
            agent_id: ID of the agent
            metadata: Metadata dictionary to merge

        Returns:
            True if updated, False if agent not found
        """
        if agent_id not in self._agents:
            return False

        self._agents[agent_id].metadata.update(metadata)
        return True

    def is_registered(self, agent_id: str) -> bool:
        """Check if an agent is registered.

        Args:
            agent_id: ID of the agent

        Returns:
            True if registered, False otherwise
        """
        return agent_id in self._agents

    def get_workspace_agents(self, workspace_id: str) -> List[AgentInfo]:
        """Get all agents in a workspace.

        Args:
            workspace_id: ID of the workspace

        Returns:
            List of AgentInfo instances
        """
        agent_ids = self._workspace_index.get(workspace_id, set())
        return [self._agents[agent_id] for agent_id in agent_ids if agent_id in self._agents]

    def get_agent_count(self, workspace_id: Optional[str] = None) -> int:
        """Get the count of registered agents.

        Args:
            workspace_id: Optional workspace ID to filter by

        Returns:
            Number of agents
        """
        if workspace_id:
            return len(self._workspace_index.get(workspace_id, set()))
        return len(self._agents)

    def clear(self) -> None:
        """Clear all registered agents."""
        self._agents.clear()
        self._workspace_index.clear()

    def get_stats(self) -> Dict[str, Any]:
        """Get registry statistics.

        Returns:
            Dictionary with registry statistics
        """
        stats = {
            "total_agents": len(self._agents),
            "workspaces": len(self._workspace_index),
            "agents_by_workspace": {
                ws_id: len(agent_ids)
                for ws_id, agent_ids in self._workspace_index.items()
            },
            "agents_by_type": {},
            "agents_by_status": {},
        }

        for agent in self._agents.values():
            # Count by type
            agent_type = agent.agent_type
            stats["agents_by_type"][agent_type] = (
                stats["agents_by_type"].get(agent_type, 0) + 1
            )

            # Count by status
            status = agent.status
            stats["agents_by_status"][status] = (
                stats["agents_by_status"].get(status, 0) + 1
            )

        return stats


# Global registry instance
_global_registry: Optional[AgentRegistry] = None


def get_registry() -> AgentRegistry:
    """Get the global agent registry instance.

    Returns:
        AgentRegistry instance
    """
    global _global_registry
    if _global_registry is None:
        _global_registry = AgentRegistry()
    return _global_registry


def reset_registry() -> None:
    """Reset the global registry (useful for testing)."""
    global _global_registry
    _global_registry = None