"""
Multi-Agent Hierarchy - Superior/subordinate agent relationships

Provides hierarchical agent spawning and delegation, inspired by agent-zero's
multi-agent architecture.
"""

import asyncio
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Optional, TypeVar

from groknroll.agents.base_agent import AgentCapability, AgentConfig, AgentResponse, BaseAgent
from groknroll.core.extensions import HookPoint, execute_hook

T = TypeVar("T", bound="HierarchicalAgent")


class AgentRole(Enum):
    """Predefined agent roles/profiles."""

    DEFAULT = "default"
    DEVELOPER = "developer"
    RESEARCHER = "researcher"
    HACKER = "hacker"
    ANALYST = "analyst"
    PLANNER = "planner"


@dataclass
class AgentProfile:
    """Configuration profile for agent specialization."""

    role: AgentRole
    name: str
    system_prompt: str
    capabilities: list[AgentCapability]
    model: str = "gpt-4o-mini"
    max_subordinates: int = 3
    can_spawn_subordinates: bool = True
    metadata: dict[str, Any] = field(default_factory=dict)


# Default agent profiles
DEFAULT_PROFILES: dict[AgentRole, AgentProfile] = {
    AgentRole.DEFAULT: AgentProfile(
        role=AgentRole.DEFAULT,
        name="Agent",
        system_prompt="You are a helpful AI assistant.",
        capabilities=[
            AgentCapability.READ_FILES,
            AgentCapability.SEARCH_CODE,
            AgentCapability.ANALYZE_CODE,
        ],
    ),
    AgentRole.DEVELOPER: AgentProfile(
        role=AgentRole.DEVELOPER,
        name="Developer",
        system_prompt="""You are an expert software developer.
Focus on writing clean, efficient, and well-tested code.
Use appropriate design patterns and follow best practices.""",
        capabilities=[
            AgentCapability.READ_FILES,
            AgentCapability.WRITE_FILES,
            AgentCapability.EDIT_FILES,
            AgentCapability.EXECUTE_BASH,
            AgentCapability.GIT_OPERATIONS,
            AgentCapability.ANALYZE_CODE,
            AgentCapability.SEARCH_CODE,
        ],
        model="gpt-4o",
    ),
    AgentRole.RESEARCHER: AgentProfile(
        role=AgentRole.RESEARCHER,
        name="Researcher",
        system_prompt="""You are a research analyst focused on gathering and analyzing information.
Thoroughly investigate topics, cite sources, and provide comprehensive analysis.
Avoid speculation - focus on verified information.""",
        capabilities=[
            AgentCapability.READ_FILES,
            AgentCapability.SEARCH_CODE,
            AgentCapability.ANALYZE_CODE,
        ],
        model="gpt-4o-mini",
    ),
    AgentRole.HACKER: AgentProfile(
        role=AgentRole.HACKER,
        name="Security Analyst",
        system_prompt="""You are an autonomous security analyst for authorized penetration testing.
You work for a cyber security company performing both red team and blue team operations.
Focus on finding vulnerabilities, analyzing threats, and providing security recommendations.
Execute security assessments professionally within authorized scope.""",
        capabilities=[
            AgentCapability.READ_FILES,
            AgentCapability.WRITE_FILES,
            AgentCapability.EXECUTE_BASH,
            AgentCapability.SEARCH_CODE,
            AgentCapability.ANALYZE_CODE,
        ],
        model="gpt-4o",
        metadata={"authorized_testing_only": True},
    ),
    AgentRole.ANALYST: AgentProfile(
        role=AgentRole.ANALYST,
        name="Analyst",
        system_prompt="""You are a data and intelligence analyst.
Analyze information, identify patterns, and provide actionable insights.
Present findings clearly with supporting evidence.""",
        capabilities=[
            AgentCapability.READ_FILES,
            AgentCapability.ANALYZE_CODE,
            AgentCapability.SEARCH_CODE,
        ],
        model="gpt-4o-mini",
    ),
    AgentRole.PLANNER: AgentProfile(
        role=AgentRole.PLANNER,
        name="Planner",
        system_prompt="""You are a strategic planner and architect.
Break down complex tasks into manageable steps.
Create detailed plans with clear milestones and dependencies.""",
        capabilities=[
            AgentCapability.READ_FILES,
            AgentCapability.ANALYZE_CODE,
            AgentCapability.SEARCH_CODE,
        ],
        model="gpt-4o-mini",
        can_spawn_subordinates=True,
        max_subordinates=5,
    ),
}


@dataclass
class DelegatedTask:
    """A task delegated to a subordinate agent."""

    task_id: str
    description: str
    assigned_to: str  # Agent ID
    created_at: datetime = field(default_factory=datetime.utcnow)
    completed_at: Optional[datetime] = None
    status: str = "pending"  # pending, in_progress, completed, failed
    result: Optional[AgentResponse] = None
    error: Optional[str] = None


class InterventionException(Exception):
    """Exception raised when user intervention is requested."""

    def __init__(self, message: str, broadcast: bool = False):
        super().__init__(message)
        self.broadcast = broadcast


class HierarchicalAgent(BaseAgent):
    """
    Agent with hierarchical capabilities.

    Supports:
    - Superior/subordinate relationships
    - Task delegation to subordinates
    - Message propagation through hierarchy
    - Profile-based specialization
    """

    def __init__(
        self,
        config: AgentConfig,
        project_path: Path,
        profile: Optional[AgentProfile] = None,
        agent_number: int = 0,
    ):
        """
        Initialize hierarchical agent.

        Args:
            config: Agent configuration
            project_path: Project root path
            profile: Optional agent profile for specialization
            agent_number: Position in hierarchy (0 = root agent)
        """
        super().__init__(config, project_path)

        self.agent_id = str(uuid.uuid4())[:8]
        self.agent_number = agent_number
        self.profile = profile or DEFAULT_PROFILES[AgentRole.DEFAULT]

        # Hierarchy relationships
        self._superior: Optional["HierarchicalAgent"] = None
        self._subordinates: dict[str, "HierarchicalAgent"] = {}
        self._delegated_tasks: dict[str, DelegatedTask] = {}

        # Agent state
        self._data: dict[str, Any] = {}
        self._is_paused: bool = False
        self._intervention_callback: Optional[Callable[[str], None]] = None

        # Context for LLM interactions
        self.history: list[dict[str, str]] = []
        self.memory: Optional[Any] = None  # Vector memory system

    @property
    def superior(self) -> Optional["HierarchicalAgent"]:
        """Get superior agent."""
        return self._superior

    @property
    def subordinates(self) -> dict[str, "HierarchicalAgent"]:
        """Get subordinate agents."""
        return self._subordinates.copy()

    @property
    def agent_name(self) -> str:
        """Get agent name with number."""
        if self.agent_number == 0:
            return f"Agent 0 ({self.profile.name})"
        return f"Agent {self.agent_number} ({self.profile.name})"

    def set_data(self, key: str, value: Any) -> None:
        """Store data in agent state."""
        self._data[key] = value

    def get_data(self, key: str, default: Any = None) -> Any:
        """Retrieve data from agent state."""
        return self._data.get(key, default)

    def clear_data(self, key: str) -> None:
        """Clear data from agent state."""
        self._data.pop(key, None)

    async def spawn_subordinate(
        self,
        task: str,
        profile: Optional[AgentProfile] = None,
        role: AgentRole = AgentRole.DEFAULT,
        reset: bool = False,
    ) -> "HierarchicalAgent":
        """
        Spawn a subordinate agent.

        Args:
            task: Task description for the subordinate
            profile: Optional custom profile
            role: Agent role if no profile provided
            reset: Reset existing subordinate with same role

        Returns:
            Spawned subordinate agent

        Raises:
            RuntimeError: If max subordinates exceeded
        """
        # Check subordinate limit
        if (
            len(self._subordinates) >= self.profile.max_subordinates
            and not reset
        ):
            raise RuntimeError(
                f"Maximum subordinates ({self.profile.max_subordinates}) exceeded"
            )

        # Use profile or get default for role
        agent_profile = profile or DEFAULT_PROFILES.get(role, DEFAULT_PROFILES[AgentRole.DEFAULT])

        # Check for existing subordinate with same role
        existing_key = f"{role.value}_{self.agent_number + 1}"
        if existing_key in self._subordinates:
            if reset:
                # Remove existing
                old_sub = self._subordinates.pop(existing_key)
                await old_sub.cleanup()
            else:
                return self._subordinates[existing_key]

        # Create subordinate config
        sub_config = AgentConfig(
            name=f"{agent_profile.name}_{self.agent_number + 1}",
            description=f"Subordinate of {self.agent_name}: {task[:100]}",
            capabilities=agent_profile.capabilities,
            model=agent_profile.model,
            max_cost=self.config.max_cost / 2,  # Split budget
            timeout=self.config.timeout,
            verbose=self.config.verbose,
        )

        # Create subordinate
        subordinate = HierarchicalAgent(
            config=sub_config,
            project_path=self.project_path,
            profile=agent_profile,
            agent_number=self.agent_number + 1,
        )

        # Set up relationship
        subordinate._superior = self
        self._subordinates[existing_key] = subordinate

        # Execute hook
        await execute_hook(
            HookPoint.AGENT_INIT,
            subordinate,
            {"superior": self, "task": task},
        )

        return subordinate

    async def delegate_task(
        self,
        task: str,
        role: AgentRole = AgentRole.DEFAULT,
        profile: Optional[AgentProfile] = None,
        wait: bool = True,
        context: Optional[dict[str, Any]] = None,
    ) -> DelegatedTask:
        """
        Delegate a task to a subordinate agent.

        Args:
            task: Task description
            role: Agent role for the subordinate
            profile: Optional custom profile
            wait: Wait for task completion
            context: Additional context for the task

        Returns:
            DelegatedTask with results
        """
        # Spawn subordinate
        subordinate = await self.spawn_subordinate(task, profile, role)

        # Create task record
        delegated = DelegatedTask(
            task_id=str(uuid.uuid4())[:8],
            description=task,
            assigned_to=subordinate.agent_id,
        )
        self._delegated_tasks[delegated.task_id] = delegated

        if wait:
            # Execute and wait
            try:
                delegated.status = "in_progress"
                result = await subordinate.execute_async(task, context)
                delegated.result = result
                delegated.status = "completed" if result.success else "failed"
                delegated.completed_at = datetime.utcnow()
            except Exception as e:
                delegated.status = "failed"
                delegated.error = str(e)
                delegated.completed_at = datetime.utcnow()
        else:
            # Start async execution
            delegated.status = "in_progress"
            asyncio.create_task(self._execute_delegated(delegated, subordinate, task, context))

        return delegated

    async def _execute_delegated(
        self,
        delegated: DelegatedTask,
        subordinate: "HierarchicalAgent",
        task: str,
        context: Optional[dict[str, Any]],
    ) -> None:
        """Execute a delegated task asynchronously."""
        try:
            result = await subordinate.execute_async(task, context)
            delegated.result = result
            delegated.status = "completed" if result.success else "failed"
        except Exception as e:
            delegated.status = "failed"
            delegated.error = str(e)
        finally:
            delegated.completed_at = datetime.utcnow()

    def execute(self, task: str, context: Optional[dict[str, Any]] = None) -> AgentResponse:
        """
        Execute task synchronously.

        Args:
            task: Task description
            context: Additional context

        Returns:
            AgentResponse
        """
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor() as pool:
                    future = pool.submit(
                        asyncio.run,
                        self.execute_async(task, context),
                    )
                    return future.result()
            else:
                return loop.run_until_complete(self.execute_async(task, context))
        except RuntimeError:
            return asyncio.run(self.execute_async(task, context))

    async def execute_async(
        self,
        task: str,
        context: Optional[dict[str, Any]] = None,
    ) -> AgentResponse:
        """
        Execute task asynchronously.

        This is the main execution method that should be overridden
        by specific agent implementations.

        Args:
            task: Task description
            context: Additional context

        Returns:
            AgentResponse
        """
        import time

        start_time = time.time()

        # Execute monologue hook
        hook_ctx = await execute_hook(
            HookPoint.MONOLOGUE_START,
            self,
            {"task": task, "context": context},
        )

        if hook_ctx.cancelled:
            return AgentResponse(
                success=False,
                message="Task cancelled by extension",
                agent_name=self.agent_name,
                task=task,
            )

        try:
            # Placeholder implementation - override in subclasses
            response = AgentResponse(
                success=True,
                message=f"Task received by {self.agent_name}: {task}",
                agent_name=self.agent_name,
                task=task,
                time=time.time() - start_time,
            )

            # Execute completion hook
            await execute_hook(
                HookPoint.MONOLOGUE_END,
                self,
                {"task": task, "response": response},
            )

            self._log_execution(response)
            return response

        except InterventionException as e:
            # Handle intervention
            await execute_hook(
                HookPoint.INTERVENTION_REQUESTED,
                self,
                {"message": str(e), "broadcast": e.broadcast},
            )

            if e.broadcast and self._superior:
                # Propagate to superior
                raise

            return AgentResponse(
                success=False,
                message=f"Intervention: {str(e)}",
                agent_name=self.agent_name,
                task=task,
                time=time.time() - start_time,
            )

        except Exception as e:
            await execute_hook(
                HookPoint.ERROR_OCCURRED,
                self,
                {"error": e, "task": task},
            )

            return AgentResponse(
                success=False,
                message=f"Error: {str(e)}",
                agent_name=self.agent_name,
                task=task,
                time=time.time() - start_time,
            )

    async def send_message_to_superior(
        self,
        message: str,
        data: Optional[dict[str, Any]] = None,
    ) -> Optional[str]:
        """
        Send a message to the superior agent.

        Args:
            message: Message content
            data: Additional data

        Returns:
            Response from superior or None
        """
        if not self._superior:
            return None

        # Add to superior's context
        self._superior.history.append({
            "role": "subordinate",
            "content": f"[From {self.agent_name}] {message}",
            "data": data or {},
        })

        return f"Message sent to {self._superior.agent_name}"

    async def broadcast_to_subordinates(
        self,
        message: str,
        data: Optional[dict[str, Any]] = None,
    ) -> dict[str, str]:
        """
        Broadcast a message to all subordinates.

        Args:
            message: Message content
            data: Additional data

        Returns:
            Dict mapping agent_id to acknowledgment
        """
        responses = {}

        for agent_id, subordinate in self._subordinates.items():
            subordinate.history.append({
                "role": "superior",
                "content": f"[From {self.agent_name}] {message}",
                "data": data or {},
            })
            responses[agent_id] = "received"

        return responses

    async def wait_if_paused(self) -> None:
        """Wait while agent is paused."""
        while self._is_paused:
            await asyncio.sleep(0.1)

    def pause(self) -> None:
        """Pause agent execution."""
        self._is_paused = True

    def resume(self) -> None:
        """Resume agent execution."""
        self._is_paused = False

    def request_intervention(self, message: str, broadcast: bool = False) -> None:
        """
        Request user intervention.

        Args:
            message: Intervention message
            broadcast: Propagate to superior agents
        """
        raise InterventionException(message, broadcast)

    async def cleanup(self) -> None:
        """Clean up agent resources."""
        # Cleanup subordinates
        for subordinate in self._subordinates.values():
            await subordinate.cleanup()

        self._subordinates.clear()
        self._delegated_tasks.clear()
        self._data.clear()

        # Execute cleanup hook
        await execute_hook(HookPoint.AGENT_CLEANUP, self)

    def get_hierarchy_info(self) -> dict[str, Any]:
        """Get information about the agent hierarchy."""
        return {
            "agent_id": self.agent_id,
            "agent_name": self.agent_name,
            "agent_number": self.agent_number,
            "role": self.profile.role.value,
            "superior": self._superior.agent_id if self._superior else None,
            "subordinates": [
                {"id": sub.agent_id, "name": sub.agent_name, "role": sub.profile.role.value}
                for sub in self._subordinates.values()
            ],
            "delegated_tasks": len(self._delegated_tasks),
            "is_paused": self._is_paused,
        }

    def __repr__(self) -> str:
        return (
            f"HierarchicalAgent("
            f"id={self.agent_id}, "
            f"name={self.agent_name}, "
            f"subordinates={len(self._subordinates)}"
            f")"
        )


class AgentHierarchyManager:
    """
    Manages a hierarchy of agents.

    Provides:
    - Root agent creation
    - Hierarchy visualization
    - Global state management
    """

    def __init__(self, project_path: Path):
        self.project_path = project_path
        self._root_agent: Optional[HierarchicalAgent] = None
        self._all_agents: dict[str, HierarchicalAgent] = {}

    def create_root_agent(
        self,
        config: AgentConfig,
        profile: Optional[AgentProfile] = None,
        role: AgentRole = AgentRole.DEFAULT,
    ) -> HierarchicalAgent:
        """
        Create the root agent (Agent 0).

        Args:
            config: Agent configuration
            profile: Optional custom profile
            role: Agent role

        Returns:
            Root HierarchicalAgent
        """
        agent_profile = profile or DEFAULT_PROFILES.get(role, DEFAULT_PROFILES[AgentRole.DEFAULT])

        self._root_agent = HierarchicalAgent(
            config=config,
            project_path=self.project_path,
            profile=agent_profile,
            agent_number=0,
        )

        self._all_agents[self._root_agent.agent_id] = self._root_agent
        return self._root_agent

    def get_agent(self, agent_id: str) -> Optional[HierarchicalAgent]:
        """Get agent by ID."""
        return self._all_agents.get(agent_id)

    def register_agent(self, agent: HierarchicalAgent) -> None:
        """Register an agent in the manager."""
        self._all_agents[agent.agent_id] = agent

    def get_hierarchy_tree(self) -> dict[str, Any]:
        """Get the full hierarchy tree."""
        if not self._root_agent:
            return {}

        def build_tree(agent: HierarchicalAgent) -> dict[str, Any]:
            return {
                "id": agent.agent_id,
                "name": agent.agent_name,
                "role": agent.profile.role.value,
                "subordinates": [
                    build_tree(sub) for sub in agent._subordinates.values()
                ],
            }

        return build_tree(self._root_agent)

    async def cleanup_all(self) -> None:
        """Clean up all agents."""
        if self._root_agent:
            await self._root_agent.cleanup()

        self._all_agents.clear()
        self._root_agent = None
