"""
Agent Spawner - Dynamic agent creation and management.

This module handles spawning specialized agents based on task requirements:
- Creates Claude Code instances with specialized prompts
- Manages agent lifecycle (spawn, monitor, terminate)
- Tracks spawned agent resources
- Integrates with the template library for rich agent definitions
"""

import asyncio
import logging
import subprocess
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Callable

from ..prompts.orchestrator_system import get_specialized_prompt, SPECIALIZED_AGENT_PROMPTS

# Import template library for rich agent definitions
try:
    from ..agents.templates import get_template, list_templates, get_registry
    from ..agents.templates.base import AgentTemplate
    TEMPLATES_AVAILABLE = True
except ImportError:
    TEMPLATES_AVAILABLE = False
    AgentTemplate = None


logger = logging.getLogger(__name__)


class AgentState(Enum):
    """State of a spawned agent."""
    INITIALIZING = "initializing"
    READY = "ready"
    BUSY = "busy"
    COMPLETED = "completed"
    FAILED = "failed"
    TERMINATED = "terminated"


@dataclass
class SpawnedAgent:
    """Represents a dynamically spawned agent."""
    id: str
    agent_type: str
    prompt: str
    workspace: str
    state: AgentState = AgentState.INITIALIZING
    spawned_at: datetime = field(default_factory=datetime.now)
    process: Optional[subprocess.Popen] = None
    tmux_session: Optional[str] = None
    tasks_completed: int = 0
    tasks_failed: int = 0
    last_activity: Optional[datetime] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


class AgentSpawner:
    """
    Manages dynamic spawning of specialized agents.

    The orchestrator can request specialized agents for specific tasks.
    This class handles:
    - Creating new Claude Code instances with custom prompts
    - Managing tmux sessions for agent isolation
    - Tracking agent lifecycle and resource usage
    """

    def __init__(
        self,
        workspace_root: Path,
        max_concurrent_agents: int = 5,
        on_agent_ready: Optional[Callable[[str], None]] = None,
        on_agent_completed: Optional[Callable[[str, str], None]] = None,
    ):
        """
        Initialize the agent spawner.

        Args:
            workspace_root: Root directory for agent workspaces
            max_concurrent_agents: Maximum number of concurrent spawned agents
            on_agent_ready: Callback when agent is ready
            on_agent_completed: Callback when agent completes (agent_id, result)
        """
        self.workspace_root = Path(workspace_root)
        self.max_concurrent_agents = max_concurrent_agents
        self.on_agent_ready = on_agent_ready
        self.on_agent_completed = on_agent_completed

        self._agents: Dict[str, SpawnedAgent] = {}
        self._lock = asyncio.Lock()

    @property
    def available_types(self) -> List[str]:
        """Get list of available agent types from templates and legacy prompts."""
        types = set(SPECIALIZED_AGENT_PROMPTS.keys())

        # Add template IDs if available
        if TEMPLATES_AVAILABLE:
            for template in list_templates():
                types.add(template.id)

        return sorted(types)

    def get_agent(self, agent_id: str) -> Optional[SpawnedAgent]:
        """Get a spawned agent by ID."""
        return self._agents.get(agent_id)

    def get_all_agents(self) -> List[SpawnedAgent]:
        """Get all spawned agents."""
        return list(self._agents.values())

    def get_active_agents(self) -> List[SpawnedAgent]:
        """Get all active (non-terminated) agents."""
        return [
            a for a in self._agents.values()
            if a.state not in (AgentState.COMPLETED, AgentState.TERMINATED, AgentState.FAILED)
        ]

    async def spawn(
        self,
        agent_type: str,
        task_context: Optional[str] = None,
        workspace: Optional[str] = None,
        custom_prompt: Optional[str] = None,
        template: Optional["AgentTemplate"] = None,
    ) -> SpawnedAgent:
        """
        Spawn a new specialized agent.

        Args:
            agent_type: Type of agent (e.g., "frontend-specialist", "story-writer")
            task_context: Optional context about the task
            workspace: Optional workspace path (defaults to workspace_root)
            custom_prompt: Optional custom prompt override
            template: Optional AgentTemplate to use directly

        Returns:
            SpawnedAgent instance

        Raises:
            ValueError: If max agents reached or invalid type
        """
        async with self._lock:
            # Check limits
            active = len(self.get_active_agents())
            if active >= self.max_concurrent_agents:
                raise ValueError(
                    f"Maximum concurrent agents ({self.max_concurrent_agents}) reached. "
                    f"Active: {active}"
                )

            # Generate agent ID
            agent_id = f"{agent_type}-{uuid.uuid4().hex[:6]}"

            # Get prompt - priority: custom_prompt > template > legacy prompt
            if custom_prompt:
                prompt = custom_prompt
            elif template:
                prompt = self._build_prompt_from_template(template, task_context)
            elif TEMPLATES_AVAILABLE:
                # Try to get template from registry
                found_template = get_template(agent_type)
                if found_template:
                    prompt = self._build_prompt_from_template(found_template, task_context)
                else:
                    # Fall back to legacy prompt
                    prompt = get_specialized_prompt(agent_type, task_context or "")
            else:
                # No templates available, use legacy
                prompt = get_specialized_prompt(agent_type, task_context or "")

            # Determine workspace
            agent_workspace = workspace or str(self.workspace_root)

            # Create agent record
            agent = SpawnedAgent(
                id=agent_id,
                agent_type=agent_type,
                prompt=prompt,
                workspace=agent_workspace,
            )

            # Store template metadata if available
            if template:
                agent.metadata["template_id"] = template.id
                agent.metadata["template_name"] = template.name
                agent.metadata["domain"] = template.domain.value

            self._agents[agent_id] = agent
            logger.info(f"Spawning agent: {agent_id} (type: {agent_type})")

            # Start the agent
            await self._start_agent(agent)

            return agent

    def _build_prompt_from_template(
        self,
        template: "AgentTemplate",
        task_context: Optional[str] = None
    ) -> str:
        """
        Build a complete prompt from an agent template.

        Args:
            template: The agent template
            task_context: Optional task-specific context

        Returns:
            Complete system prompt string
        """
        # Start with the template's system prompt
        prompt_parts = [template.system_prompt]

        # Add capabilities section
        if template.capabilities:
            cap_list = ", ".join(c.name for c in template.capabilities)
            prompt_parts.append(f"\n\nYour capabilities: {cap_list}")

        # Add suggested tools
        if template.suggested_tools:
            tools_list = ", ".join(template.suggested_tools)
            prompt_parts.append(f"\nSuggested tools: {tools_list}")

        # Add collaboration style
        prompt_parts.append(f"\nCollaboration style: {template.collaboration_style.value}")

        # Add task context if provided
        if task_context:
            prompt_parts.append(f"\n\n## Current Task Context\n{task_context}")

        return "\n".join(prompt_parts)

    async def _start_agent(self, agent: SpawnedAgent) -> None:
        """
        Start a spawned agent in a tmux session.

        This creates a new tmux session with Claude Code initialized
        with the specialized prompt.
        """
        # Create tmux session name
        session_name = f"agent-{agent.id}"
        agent.tmux_session = session_name

        try:
            # Create prompt file for the agent
            prompt_file = self.workspace_root / f".agent-{agent.id}-prompt.md"
            prompt_file.write_text(agent.prompt)

            # Check if tmux is available
            result = subprocess.run(
                ["which", "tmux"],
                capture_output=True,
                text=True
            )

            if result.returncode != 0:
                # Tmux not available, mark as ready anyway (for testing)
                logger.warning("Tmux not available, running in simulated mode")
                agent.state = AgentState.READY
                agent.last_activity = datetime.now()
                return

            # Create new tmux session
            subprocess.run(
                [
                    "tmux", "new-session",
                    "-d",  # Detached
                    "-s", session_name,
                    "-c", agent.workspace,  # Working directory
                ],
                check=True
            )

            # Send command to start Claude Code with the prompt
            # The agent will read the prompt file on startup
            claude_cmd = f"claude --prompt-file {prompt_file}"
            subprocess.run(
                [
                    "tmux", "send-keys",
                    "-t", session_name,
                    claude_cmd,
                    "Enter"
                ],
                check=True
            )

            agent.state = AgentState.READY
            agent.last_activity = datetime.now()

            logger.info(f"Agent {agent.id} started in tmux session: {session_name}")

            if self.on_agent_ready:
                self.on_agent_ready(agent.id)

        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to start agent {agent.id}: {e}")
            agent.state = AgentState.FAILED
            raise

    async def send_task(self, agent_id: str, task: str) -> bool:
        """
        Send a task to a spawned agent.

        Args:
            agent_id: The agent ID
            task: Task description to send

        Returns:
            True if task was sent successfully
        """
        agent = self._agents.get(agent_id)
        if not agent:
            logger.error(f"Agent not found: {agent_id}")
            return False

        if agent.state != AgentState.READY:
            logger.warning(f"Agent {agent_id} not ready (state: {agent.state})")
            return False

        if not agent.tmux_session:
            logger.error(f"Agent {agent_id} has no tmux session")
            return False

        try:
            # Send the task to the agent's tmux session
            subprocess.run(
                [
                    "tmux", "send-keys",
                    "-t", agent.tmux_session,
                    task,
                    "Enter"
                ],
                check=True
            )

            agent.state = AgentState.BUSY
            agent.last_activity = datetime.now()

            logger.info(f"Task sent to agent {agent_id}")
            return True

        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to send task to {agent_id}: {e}")
            return False

    async def check_status(self, agent_id: str) -> Dict[str, Any]:
        """
        Check the status of a spawned agent.

        Args:
            agent_id: The agent ID

        Returns:
            Status dictionary
        """
        agent = self._agents.get(agent_id)
        if not agent:
            return {"error": f"Agent not found: {agent_id}"}

        status = {
            "id": agent.id,
            "type": agent.agent_type,
            "state": agent.state.value,
            "workspace": agent.workspace,
            "spawned_at": agent.spawned_at.isoformat(),
            "tasks_completed": agent.tasks_completed,
            "tasks_failed": agent.tasks_failed,
            "last_activity": agent.last_activity.isoformat() if agent.last_activity else None,
        }

        # Check if tmux session is still alive
        if agent.tmux_session:
            try:
                result = subprocess.run(
                    ["tmux", "has-session", "-t", agent.tmux_session],
                    capture_output=True
                )
                status["session_alive"] = result.returncode == 0
            except Exception:
                status["session_alive"] = False

        return status

    async def terminate(self, agent_id: str, force: bool = False) -> bool:
        """
        Terminate a spawned agent.

        Args:
            agent_id: The agent ID
            force: Force terminate even if busy

        Returns:
            True if terminated successfully
        """
        agent = self._agents.get(agent_id)
        if not agent:
            logger.warning(f"Agent not found for termination: {agent_id}")
            return False

        if agent.state == AgentState.BUSY and not force:
            logger.warning(f"Agent {agent_id} is busy, use force=True to terminate")
            return False

        try:
            # Kill tmux session
            if agent.tmux_session:
                subprocess.run(
                    ["tmux", "kill-session", "-t", agent.tmux_session],
                    capture_output=True
                )

            # Clean up prompt file
            prompt_file = self.workspace_root / f".agent-{agent.id}-prompt.md"
            if prompt_file.exists():
                prompt_file.unlink()

            agent.state = AgentState.TERMINATED
            logger.info(f"Agent {agent_id} terminated")

            return True

        except Exception as e:
            logger.error(f"Error terminating agent {agent_id}: {e}")
            return False

    async def terminate_all(self, force: bool = False) -> int:
        """
        Terminate all spawned agents.

        Args:
            force: Force terminate even if busy

        Returns:
            Number of agents terminated
        """
        terminated = 0
        for agent_id in list(self._agents.keys()):
            if await self.terminate(agent_id, force=force):
                terminated += 1
        return terminated

    def mark_task_completed(self, agent_id: str, success: bool = True) -> None:
        """
        Mark a task as completed for an agent.

        Args:
            agent_id: The agent ID
            success: Whether the task succeeded
        """
        agent = self._agents.get(agent_id)
        if not agent:
            return

        if success:
            agent.tasks_completed += 1
        else:
            agent.tasks_failed += 1

        agent.state = AgentState.READY
        agent.last_activity = datetime.now()

        if self.on_agent_completed:
            self.on_agent_completed(agent_id, "success" if success else "failed")

    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about spawned agents."""
        agents = list(self._agents.values())

        by_state = {}
        for agent in agents:
            state = agent.state.value
            by_state[state] = by_state.get(state, 0) + 1

        by_type = {}
        for agent in agents:
            by_type[agent.agent_type] = by_type.get(agent.agent_type, 0) + 1

        total_completed = sum(a.tasks_completed for a in agents)
        total_failed = sum(a.tasks_failed for a in agents)

        return {
            "total_agents": len(agents),
            "active_agents": len(self.get_active_agents()),
            "max_concurrent": self.max_concurrent_agents,
            "by_state": by_state,
            "by_type": by_type,
            "tasks_completed": total_completed,
            "tasks_failed": total_failed,
            "success_rate": total_completed / (total_completed + total_failed) if (total_completed + total_failed) > 0 else 1.0,
        }
