"""
Orchestration Brain - The central intelligence for the multi-agent system.

This module is the "brain" that the Claude Code orchestrator uses to:
- Understand the current system state
- Make decisions about task assignment
- Monitor agent workload and health
- Dynamically adapt to changing conditions
"""

import asyncio
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
from pathlib import Path

from ..persistence.database import OrchestratorDB
from ..memory.operational import OperationalMemory
from ..tracking.cli_usage import CLIUsageTracker, get_cli_tracker, AgentAvailability
from ..prompts.orchestrator_system import build_orchestrator_prompt
from .commands import OrchestratorCommands, CommandResult
from .spawner import AgentSpawner, SpawnedAgent

# Import intelligence modules for smart agent selection
try:
    from ..intelligence.domain_analyzer import analyze_task, TaskDomain
    from ..agents.templates import select_agents_for_task, get_template
    SMART_ROUTING_AVAILABLE = True
except ImportError:
    SMART_ROUTING_AVAILABLE = False


logger = logging.getLogger(__name__)


class OrchestrationBrain:
    """
    The orchestration brain coordinates the multi-agent system.

    This class provides:
    - System state awareness
    - Task assignment intelligence
    - Agent workload balancing
    - Dynamic adaptation
    """

    def __init__(
        self,
        db: OrchestratorDB,
        ops_path: Path,
        workspace_root: Optional[Path] = None,
    ):
        """
        Initialize the orchestration brain.

        Args:
            db: Database for persistence
            ops_path: Path to operational files
            workspace_root: Root for agent workspaces
        """
        self.db = db
        self.ops_path = Path(ops_path)
        self.workspace_root = workspace_root or Path.cwd()

        # Initialize components
        self.operational_memory = OperationalMemory(db, ops_path)
        self.cli_tracker = get_cli_tracker()
        self.commands = OrchestratorCommands(
            db=db,
            operational_memory=self.operational_memory,
            cli_tracker=self.cli_tracker,
        )
        self.spawner = AgentSpawner(
            workspace_root=self.workspace_root,
            on_agent_ready=self._on_agent_ready,
            on_agent_completed=self._on_agent_completed,
        )

        # State
        self._adapters: Dict[str, Any] = {}
        self._last_status_check: Optional[datetime] = None
        self._warnings: List[str] = []

    def register_adapter(self, agent_id: str, adapter: Any) -> None:
        """Register an agent adapter."""
        self._adapters[agent_id] = adapter
        self.commands.register_adapter(agent_id, adapter)
        logger.info(f"Brain registered adapter: {agent_id}")

    def _on_agent_ready(self, agent_id: str) -> None:
        """Callback when a spawned agent is ready."""
        logger.info(f"Spawned agent ready: {agent_id}")

    def _on_agent_completed(self, agent_id: str, result: str) -> None:
        """Callback when a spawned agent completes a task."""
        logger.info(f"Spawned agent {agent_id} completed: {result}")

    async def process_command(self, command_str: str) -> CommandResult:
        """
        Process an orchestrator command.

        Args:
            command_str: The command (e.g., "/assign claude-code Fix bug")

        Returns:
            CommandResult with outcome
        """
        return await self.commands.execute(command_str)

    def build_system_prompt(self) -> str:
        """
        Build the complete system prompt for the orchestration agent.

        This prompt contains all context needed for the orchestrator
        to make informed decisions.
        """
        # Gather current state
        project_state = self.operational_memory.read_project_state()

        # Build available agents list
        available_agents = []
        for agent_id, adapter in self._adapters.items():
            availability = self.cli_tracker.get_availability(agent_id)
            usage = self.cli_tracker.get_usage_stats(agent_id)

            available_agents.append({
                "id": agent_id,
                "type": getattr(adapter, "agent_type", "unknown"),
                "status": availability.value,
                "capabilities": getattr(adapter, "capabilities", []),
            })

        # Add spawned agents
        for agent in self.spawner.get_active_agents():
            available_agents.append({
                "id": agent.id,
                "type": agent.agent_type,
                "status": agent.state.value,
                "capabilities": ["specialized", agent.agent_type],
            })

        # Get pending tasks
        pending_tasks = []
        for task in self.db.get_pending_tasks(limit=20):
            pending_tasks.append({
                "id": task.id,
                "description": task.description,
                "status": task.status,
                "assigned_to": task.assigned_to,
            })

        # Build memory summary
        decisions = self.operational_memory.get_active_decisions()
        memory_summary = {
            "decisions_count": len(decisions),
            "runbooks_count": 0,  # Would query knowledge layer
            "patterns_count": 0,  # Would query knowledge layer
            "recent_items": [d["decision"][:50] for d in decisions[:3]],
        }

        # Build usage stats
        usage_stats = {}
        for agent_id in self._adapters:
            usage = self.cli_tracker.get_usage_stats(agent_id)
            usage_stats[agent_id] = {
                "tokens_used": usage.get("total_tokens", 0),
                "tokens_limit": usage.get("token_limit", "session-based"),
                "session_percentage": usage.get("session_percentage", 0),
            }

        return build_orchestrator_prompt(
            project_state=project_state,
            available_agents=available_agents,
            pending_tasks=pending_tasks,
            memory_summary=memory_summary,
            usage_stats=usage_stats,
        )

    def get_system_status(self) -> Dict[str, Any]:
        """
        Get comprehensive system status.

        Returns a dictionary with all relevant system information.
        """
        self._last_status_check = datetime.now()
        self._warnings = []

        # Project state
        project_state = self.operational_memory.read_project_state()

        # Agent statuses
        agents = []
        for agent_id, adapter in self._adapters.items():
            availability = self.cli_tracker.get_availability(agent_id)
            usage = self.cli_tracker.get_usage_stats(agent_id)

            agent_info = {
                "id": agent_id,
                "type": getattr(adapter, "agent_type", "cli"),
                "availability": availability.value,
                "session_percentage": usage.get("session_percentage", 0),
                "total_requests": usage.get("total_requests", 0),
                "error_count": usage.get("error_count", 0),
            }
            agents.append(agent_info)

            # Generate warnings
            if usage.get("session_percentage", 0) > 80:
                self._warnings.append(
                    f"Agent {agent_id} at {usage['session_percentage']:.0f}% session capacity"
                )
            if availability == AgentAvailability.RATE_LIMITED:
                self._warnings.append(f"Agent {agent_id} is rate limited")

        # Spawned agents
        spawned = []
        for agent in self.spawner.get_all_agents():
            spawned.append({
                "id": agent.id,
                "type": agent.agent_type,
                "state": agent.state.value,
                "tasks_completed": agent.tasks_completed,
            })

        # Tasks
        pending_tasks = self.db.get_pending_tasks(limit=100)
        task_counts = {
            "pending": len([t for t in pending_tasks if t.status == "pending"]),
            "in_progress": len([t for t in pending_tasks if t.status == "in_progress"]),
        }

        return {
            "timestamp": datetime.now().isoformat(),
            "project": {
                "name": project_state.get("project", {}).get("name", "Unknown"),
                "phase": project_state.get("current_phase", "Unknown"),
                "version": project_state.get("version", "0.0.0"),
            },
            "agents": {
                "registered": agents,
                "spawned": spawned,
                "total_active": len(agents) + len(self.spawner.get_active_agents()),
            },
            "tasks": task_counts,
            "warnings": self._warnings,
            "spawner_stats": self.spawner.get_stats(),
        }

    def suggest_agent_for_task(
        self,
        task_description: str,
        task_type: Optional[str] = None,
    ) -> Tuple[str, str]:
        """
        Suggest the best agent for a task using intelligent domain analysis.

        Uses the domain analyzer to understand the task semantically and
        match it with the most appropriate agent template.

        Args:
            task_description: Description of the task
            task_type: Optional task type hint

        Returns:
            Tuple of (agent_id, reason)
        """
        # Use smart routing if available
        if SMART_ROUTING_AVAILABLE:
            return self._suggest_agent_smart(task_description)

        # Fallback to heuristic-based assignment
        return self._suggest_agent_heuristic(task_description)

    def _suggest_agent_smart(
        self,
        task_description: str,
    ) -> Tuple[str, str]:
        """
        Use domain analyzer and template registry for smart agent selection.

        Args:
            task_description: Description of the task

        Returns:
            Tuple of (agent_id, reason)
        """
        # Analyze the task
        analysis = analyze_task(task_description)
        domain_name = analysis.primary_domain.value.replace("_", " ").title()
        confidence = analysis.primary_confidence

        logger.info(
            f"Task analysis: domain={analysis.primary_domain.value}, "
            f"confidence={confidence:.0%}, complexity={analysis.complexity}"
        )

        # Get recommended agents from template registry
        matches = select_agents_for_task(analysis)

        if matches:
            best_match = matches[0]
            template_id = best_match.template.id
            template_name = best_match.template.name

            # Check if we already have this type of agent running
            for agent in self.spawner.get_active_agents():
                if agent.agent_type == template_id:
                    return (
                        agent.id,
                        f"{template_name} available "
                        f"({domain_name} task, {confidence:.0%} confidence)"
                    )

            # Suggest spawning the recommended agent
            return (
                f"spawn:{template_id}",
                f"Recommend spawning {template_name} "
                f"({domain_name} task, {confidence:.0%} confidence)"
            )

        # No template match - use best available CLI agent
        cli_agents = [aid for aid in self._adapters.keys() if "cli" in aid or "code" in aid]
        best = self.cli_tracker.get_best_available_agent(cli_agents)

        if best:
            usage = self.cli_tracker.get_usage_stats(best)
            return (
                best,
                f"Best available CLI agent for {domain_name} task "
                f"({usage.get('session_percentage', 0):.0f}% used)"
            )

        return "", "No suitable agent available"

    def _suggest_agent_heuristic(
        self,
        task_description: str,
    ) -> Tuple[str, str]:
        """
        Fallback heuristic-based agent selection.

        Args:
            task_description: Description of the task

        Returns:
            Tuple of (agent_id, reason)
        """
        description_lower = task_description.lower()

        # Check for specialized needs
        if any(kw in description_lower for kw in ["frontend", "ui", "css", "react", "vue"]):
            # Check if we have a frontend specialist
            for agent in self.spawner.get_active_agents():
                if agent.agent_type == "frontend-specialist":
                    return agent.id, "Specialized frontend agent available"

            # Suggest spawning one
            return "spawn:frontend-specialist", "Consider spawning a frontend specialist"

        if any(kw in description_lower for kw in ["api", "backend", "database", "server"]):
            for agent in self.spawner.get_active_agents():
                if agent.agent_type == "backend-specialist":
                    return agent.id, "Specialized backend agent available"
            return "spawn:backend-specialist", "Consider spawning a backend specialist"

        if any(kw in description_lower for kw in ["test", "spec", "coverage"]):
            for agent in self.spawner.get_active_agents():
                if agent.agent_type == "test-engineer":
                    return agent.id, "Specialized test agent available"
            return "spawn:test-engineer", "Consider spawning a test engineer"

        # Creative writing tasks
        if any(kw in description_lower for kw in ["story", "poem", "write", "creative", "fiction"]):
            for agent in self.spawner.get_active_agents():
                if agent.agent_type == "story-writer":
                    return agent.id, "Specialized story writer available"
            return "spawn:story-writer", "Consider spawning a story writer"

        # Use best available CLI agent
        cli_agents = [aid for aid in self._adapters.keys() if "cli" in aid or "code" in aid]
        best = self.cli_tracker.get_best_available_agent(cli_agents)

        if best:
            usage = self.cli_tracker.get_usage_stats(best)
            return best, f"Best available CLI agent ({usage.get('session_percentage', 0):.0f}% used)"

        # Fall back to any available
        for agent_id in self._adapters:
            availability = self.cli_tracker.get_availability(agent_id)
            if availability == AgentAvailability.AVAILABLE:
                return agent_id, "Available agent"

        return "", "No agents available"

    async def analyze_workload(self) -> Dict[str, Any]:
        """
        Analyze current workload and provide recommendations.
        """
        status = self.get_system_status()
        recommendations = []

        # Check for overloaded agents
        for agent in status["agents"]["registered"]:
            if agent["session_percentage"] > 80:
                recommendations.append({
                    "type": "rebalance",
                    "message": f"Consider redistributing work from {agent['id']}",
                    "agent_id": agent["id"],
                })

        # Check for stuck tasks (no agent assigned)
        pending = self.db.get_pending_tasks(limit=50)
        unassigned = [t for t in pending if not t.assigned_to and t.status == "pending"]
        if unassigned:
            recommendations.append({
                "type": "assign",
                "message": f"{len(unassigned)} tasks need assignment",
                "task_ids": [t.id for t in unassigned[:5]],
            })

        # Check if we should spawn specialists
        if status["tasks"]["pending"] > 5 and len(status["agents"]["spawned"]) == 0:
            recommendations.append({
                "type": "spawn",
                "message": "High task load - consider spawning specialist agents",
            })

        return {
            "timestamp": datetime.now().isoformat(),
            "workload": {
                "pending_tasks": status["tasks"]["pending"],
                "in_progress_tasks": status["tasks"]["in_progress"],
                "active_agents": status["agents"]["total_active"],
            },
            "recommendations": recommendations,
            "warnings": status["warnings"],
        }

    async def auto_assign_pending_tasks(self) -> List[Dict[str, Any]]:
        """
        Automatically assign pending tasks to available agents.

        Returns:
            List of assignments made
        """
        assignments = []
        pending = self.db.get_pending_tasks(limit=10)

        for task in pending:
            if task.assigned_to or task.status != "pending":
                continue

            agent_id, reason = self.suggest_agent_for_task(task.description)

            if agent_id and not agent_id.startswith("spawn:"):
                # Execute assignment
                result = await self.commands.execute(
                    f"/assign {agent_id} {task.description}"
                )

                if result.success:
                    assignments.append({
                        "task_id": task.id,
                        "agent_id": agent_id,
                        "reason": reason,
                    })

        return assignments

    def get_orchestrator_context(self) -> str:
        """
        Get a formatted context string for the orchestrator to display on startup.
        """
        status = self.get_system_status()

        lines = [
            "=" * 60,
            "ORCHESTRATION AGENT INITIALIZED",
            "=" * 60,
            "",
            f"Project: {status['project']['name']} v{status['project']['version']}",
            f"Phase: {status['project']['phase']}",
            "",
            "REGISTERED AGENTS:",
        ]

        for agent in status["agents"]["registered"]:
            emoji = {
                "available": "🟢",
                "busy": "🟡",
                "limited": "🟠",
                "unavailable": "🔴",
            }.get(agent["availability"], "⚪")

            lines.append(
                f"  {emoji} {agent['id']:<15} "
                f"[{agent['availability']:<12}] "
                f"Usage: {agent['session_percentage']:.0f}%"
            )

        if status["agents"]["spawned"]:
            lines.append("")
            lines.append("SPAWNED SPECIALISTS:")
            for agent in status["agents"]["spawned"]:
                lines.append(f"  • {agent['id']} ({agent['type']}) - {agent['state']}")

        lines.extend([
            "",
            f"TASKS: {status['tasks']['pending']} pending, {status['tasks']['in_progress']} in progress",
        ])

        if status["warnings"]:
            lines.extend(["", "WARNINGS:"])
            for warning in status["warnings"]:
                lines.append(f"  ⚠️  {warning}")

        lines.extend([
            "",
            "=" * 60,
            "Type /help for available commands",
            "Enter tasks to assign work to agents",
            "=" * 60,
        ])

        return "\n".join(lines)

    async def shutdown(self) -> None:
        """Clean up resources on shutdown."""
        # Terminate all spawned agents
        terminated = await self.spawner.terminate_all(force=True)
        logger.info(f"Terminated {terminated} spawned agents")
