"""Agent Manager — orchestrates all 5 persistent AI agents.

Manages the lifecycle of all agents:
- Starts/stops all agents
- Assigns goals to the best-suited agent
- Monitors agent progress
- Handles agent failures and reassignment
- Provides aggregated stats

The agents run persistently in background threads, continuously
picking up goals from the goal memory and working on them.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable

from ..memory.goal_memory import GoalMemory, Goal
from ..memory.persistent import PersistentMemory
from .agent_base import BaseAgent
from .planner_agent import PlannerAgent
from .coder_agent import CoderAgent
from .researcher_agent import ResearcherAgent
from .reviewer_agent import ReviewerAgent
from .executor_agent import ExecutorAgent

logger = logging.getLogger(__name__)


class AgentManager:
    """Manages 5 persistent AI agents that work on goals and projects.

    Agents:
    1. Planner — breaks goals into steps
    2. Coder — writes code, creates files
    3. Researcher — gathers information
    4. Reviewer — validates quality
    5. Executor — runs commands, executes tools

    All agents run in background threads and continuously process goals.
    """

    def __init__(self, goal_memory: GoalMemory,
                 persistent_memory: PersistentMemory | None = None,
                 generate_fn: Callable[[str], str] | None = None,
                 tool_registry: Any = None) -> None:
        self.goal_memory = goal_memory
        self.persistent_memory = persistent_memory
        self._generate_fn = generate_fn
        self._tool_registry = tool_registry

        # Create all 5 agents
        self.agents: dict[str, BaseAgent] = {
            "planner": PlannerAgent(goal_memory, persistent_memory, generate_fn),
            "coder": CoderAgent(goal_memory, persistent_memory, generate_fn),
            "researcher": ResearcherAgent(goal_memory, persistent_memory, generate_fn),
            "reviewer": ReviewerAgent(goal_memory, persistent_memory, generate_fn),
            "executor": ExecutorAgent(goal_memory, persistent_memory, generate_fn, tool_registry),
        }

        self._running = False
        self._start_time = 0.0
        self._stats = {
            "total_goals_assigned": 0,
            "total_goals_completed": 0,
            "total_goals_failed": 0,
            "reassignments": 0,
        }

    def set_generate_fn(self, fn: Callable[[str], str]) -> None:
        """Set the LLM generation function for all agents."""
        self._generate_fn = fn
        for agent in self.agents.values():
            agent.set_generate_fn(fn)

    def start_all(self) -> None:
        """Start all 5 agents."""
        self._running = True
        self._start_time = time.time()
        for agent in self.agents.values():
            agent.start()
        logger.info("All 5 agents started")

    def stop_all(self) -> None:
        """Stop all agents."""
        self._running = False
        for agent in self.agents.values():
            agent.stop()
        logger.info("All agents stopped")

    def pause_all(self) -> None:
        """Pause all agents."""
        for agent in self.agents.values():
            agent.pause()

    def resume_all(self) -> None:
        """Resume all agents."""
        for agent in self.agents.values():
            agent.resume()

    def create_project(self, title: str, description: str, priority: str = "high",
                       tags: list[str] | None = None) -> Goal:
        """Create a new project goal and let the planner agent pick it up."""
        goal = self.goal_memory.create_goal(
            title=title, description=description, priority=priority,
            tags=tags or ["project"],
        )
        self._stats["total_goals_assigned"] += 1
        logger.info("Created project: %s (will be picked up by planner agent)", title)
        return goal

    def assign_goal(self, goal_id: str, agent_name: str) -> bool:
        """Manually assign a goal to a specific agent."""
        if agent_name not in self.agents:
            logger.warning("Unknown agent: %s", agent_name)
            return False
        result = self.goal_memory.assign_agent(goal_id, agent_name)
        if result:
            self._stats["total_goals_assigned"] += 1
            return True
        return False

    def get_active_projects(self) -> list[Goal]:
        """Get all active goals/projects."""
        return self.goal_memory.get_active_goals()

    def get_agent_status(self) -> dict[str, Any]:
        """Get status of all agents."""
        return {name: agent.get_status() for name, agent in self.agents.items()}

    def get_project_progress(self) -> list[dict[str, Any]]:
        """Get progress of all active projects."""
        active = self.goal_memory.get_active_goals()
        return [g.as_dict() for g in active]

    def get_stats(self) -> dict[str, Any]:
        """Get aggregated stats."""
        agent_stats = {name: agent.get_status()["stats"] for name, agent in self.agents.items()}
        goal_stats = self.goal_memory.get_stats()

        total_completed = sum(s["goals_completed"] for s in agent_stats.values())
        total_failed = sum(s["goals_failed"] for s in agent_stats.values())
        total_steps = sum(s["steps_executed"] for s in agent_stats.values())

        return {
            "manager": {
                **self._stats,
                "running": self._running,
                "uptime_s": round(time.time() - self._start_time, 1) if self._start_time else 0,
                "total_goals_completed": total_completed,
                "total_goals_failed": total_failed,
                "total_steps_executed": total_steps,
            },
            "agents": agent_stats,
            "goals": goal_stats,
        }