"""
Agent Launcher

Launches specialized agents for parallel or sequential execution.
"""

import asyncio
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn

from ai_coder.agents.base import Agent, AgentResult, AgentType
from ai_coder.agents.code_explorer import CodeExplorerAgent
from ai_coder.agents.code_architect import CodeArchitectAgent
from ai_coder.agents.code_reviewer import CodeReviewerAgent
from ai_coder.agents.planner import PlannerAgent
from ai_coder.agents.security_reviewer import SecurityReviewerAgent
from ai_coder.agents.tdd_guide import TDDGuideAgent
from ai_coder.agents.build_fixer import BuildFixerAgent
from ai_coder.agents.debugger import DebuggerAgent
from ai_coder.agents.refactor import RefactorAgent
from ai_coder.llm.interface import LLMProvider
from ai_coder.tools.base import ToolRegistry

console = Console()


@dataclass
class AgentTask:
    """A task for an agent to execute."""
    agent_type: AgentType
    prompt: str


class AgentLauncher:
    """
    Launches specialized agents for parallel or sequential execution.
    
    Supports:
    - All 10 agent types
    - Parallel execution of multiple agents
    - Sequential execution
    """
    
    def __init__(
        self,
        llm: LLMProvider,
        tool_registry: ToolRegistry,
        project_root: Path,
        max_workers: int = 3,
    ):
        self.llm = llm
        self.tool_registry = tool_registry
        self.project_root = project_root
        self.max_workers = max_workers
        
        # Agent class mapping - all 10 types
        self._agent_classes = {
            AgentType.CODE_EXPLORER: CodeExplorerAgent,
            AgentType.CODE_ARCHITECT: CodeArchitectAgent,
            AgentType.CODE_REVIEWER: CodeReviewerAgent,
            AgentType.PLANNER: PlannerAgent,
            AgentType.SECURITY_REVIEWER: SecurityReviewerAgent,
            AgentType.TDD_GUIDE: TDDGuideAgent,
            AgentType.BUILD_FIXER: BuildFixerAgent,
            AgentType.DEBUGGER: DebuggerAgent,
            AgentType.REFACTOR: RefactorAgent,
        }

    def _create_agent(self, agent_type: AgentType) -> Agent:
        """Create an agent instance by type."""
        # specific agent implementation or generic DefinedAgent?
        
        # Check for definition file in package
        # Assuming __file__ is inside ai_coder/agents/launcher.py
        package_root = Path(__file__).parent.parent
        def_name = agent_type.value # e.g. "build-fixer"
        def_path = package_root / "agents" / "definitions" / f"{def_name}.md"
        
        if def_path.exists():
            from ai_coder.agents.defined import DefinedAgent
            return DefinedAgent(
                llm=self.llm,
                tool_registry=self.tool_registry,
                project_root=self.project_root,
                definition_path=def_path,
                agent_type=agent_type
            )

        agent_class = self._agent_classes.get(agent_type)
        if not agent_class:
            raise ValueError(f"Unknown agent type: {agent_type}")
        
        return agent_class(
            llm=self.llm,
            tool_registry=self.tool_registry,
            project_root=self.project_root,
        )

    def launch_single(self, agent_type: AgentType, prompt: str) -> AgentResult:
        """Launch a single agent with the given prompt."""
        agent = self._create_agent(agent_type)
        return agent.execute(prompt)

    def launch_parallel(
        self,
        tasks: list[AgentTask],
    ) -> list[AgentResult]:
        """Launch multiple agents in parallel."""
        results = []
        
        console.print(f"[bold cyan]Launching {len(tasks)} agents in parallel...[/bold cyan]")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
        ) as progress:
            task = progress.add_task("Running agents...", total=len(tasks))
            
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                future_to_task = {
                    executor.submit(
                        self._execute_agent_task, t
                    ): t for t in tasks
                }
                
                for future in as_completed(future_to_task):
                    agent_task = future_to_task[future]
                    try:
                        result = future.result()
                        results.append(result)
                        progress.update(task, advance=1)
                        console.print(
                            f"[green]✓ {agent_task.agent_type.display_name} complete[/green]"
                        )
                    except Exception as e:
                        results.append(
                            AgentResult.error_result(agent_task.agent_type, str(e))
                        )
                        progress.update(task, advance=1)
                        console.print(
                            f"[red]✗ {agent_task.agent_type.display_name} failed: {e}[/red]"
                        )
        
        return results

    def _execute_agent_task(self, task: AgentTask) -> AgentResult:
        """Execute a single agent task."""
        agent = self._create_agent(task.agent_type)
        return agent.execute(task.prompt)

    def launch_explorers(self, prompts: list[str]) -> list[AgentResult]:
        """Launch multiple code-explorer agents in parallel."""
        tasks = [AgentTask(AgentType.CODE_EXPLORER, p) for p in prompts]
        return self.launch_parallel(tasks)

    def launch_architects(self, prompts: list[str]) -> list[AgentResult]:
        """Launch multiple code-architect agents in parallel."""
        tasks = [AgentTask(AgentType.CODE_ARCHITECT, p) for p in prompts]
        return self.launch_parallel(tasks)

    def launch_reviewers(self, prompts: list[str]) -> list[AgentResult]:
        """Launch multiple code-reviewer agents in parallel."""
        tasks = [AgentTask(AgentType.CODE_REVIEWER, p) for p in prompts]
        return self.launch_parallel(tasks)
