"""Agent orchestrator for multi-agent collaboration.

This module provides the main orchestration logic for coordinating
multiple specialized agents working on complex tasks.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional, Any, TYPE_CHECKING
from enum import Enum

from ieum.agents.roles import AgentRole, AgentConfig, AGENT_CONFIGS, get_agents_for_task

if TYPE_CHECKING:
    from langchain_core.language_models import BaseChatModel


class TaskStatus(Enum):
    """Status of a task."""

    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    BLOCKED = "blocked"


@dataclass
class AgentResult:
    """Result from an agent execution."""

    agent_role: AgentRole
    task_id: str
    success: bool
    output: str
    artifacts: list[str] = field(default_factory=list)
    next_tasks: list[str] = field(default_factory=list)
    error: Optional[str] = None
    duration_seconds: float = 0.0
    timestamp: datetime = field(default_factory=datetime.now)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "agent_role": self.agent_role.value,
            "task_id": self.task_id,
            "success": self.success,
            "output": self.output,
            "artifacts": self.artifacts,
            "next_tasks": self.next_tasks,
            "error": self.error,
            "duration_seconds": self.duration_seconds,
            "timestamp": self.timestamp.isoformat(),
        }


@dataclass
class OrchestratorTask:
    """Task for the orchestrator."""

    id: str
    description: str
    assigned_to: Optional[AgentRole] = None
    status: TaskStatus = TaskStatus.PENDING
    dependencies: list[str] = field(default_factory=list)
    result: Optional[AgentResult] = None
    created_at: datetime = field(default_factory=datetime.now)
    completed_at: Optional[datetime] = None


@dataclass
class WorkflowResult:
    """Result of a complete workflow execution."""

    success: bool
    results: list[AgentResult] = field(default_factory=list)
    final_output: str = ""
    total_duration: float = 0.0
    errors: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "success": self.success,
            "results": [r.to_dict() for r in self.results],
            "final_output": self.final_output,
            "total_duration": self.total_duration,
            "errors": self.errors,
        }


class AgentOrchestrator:
    """Orchestrates multiple agents for complex tasks."""

    def __init__(self, model: Optional["BaseChatModel"] = None):
        """Initialize the orchestrator.

        Args:
            model: LLM model for agent execution.
        """
        self.model = model
        self.tasks: dict[str, OrchestratorTask] = {}
        self.results: list[AgentResult] = []
        self.shared_context: dict[str, Any] = {}
        self._task_counter = 0

    def create_task(
        self,
        description: str,
        assigned_to: Optional[AgentRole] = None,
        dependencies: Optional[list[str]] = None,
    ) -> str:
        """Create a new task.

        Args:
            description: Task description.
            assigned_to: Optional agent role.
            dependencies: Optional task dependencies.

        Returns:
            Task ID.
        """
        self._task_counter += 1
        task_id = f"task_{self._task_counter}"

        self.tasks[task_id] = OrchestratorTask(
            id=task_id,
            description=description,
            assigned_to=assigned_to,
            dependencies=dependencies or [],
        )

        return task_id

    def assign_task(self, task_id: str, role: AgentRole) -> bool:
        """Assign a task to an agent role.

        Args:
            task_id: Task ID.
            role: Agent role to assign.

        Returns:
            True if assignment successful.
        """
        if task_id not in self.tasks:
            return False

        self.tasks[task_id].assigned_to = role
        return True

    async def execute_task(self, task_id: str) -> AgentResult:
        """Execute a single task.

        Args:
            task_id: Task ID to execute.

        Returns:
            AgentResult from execution.
        """
        if task_id not in self.tasks:
            return AgentResult(
                agent_role=AgentRole.IMPLEMENTER,
                task_id=task_id,
                success=False,
                output="",
                error="Task not found",
            )

        task = self.tasks[task_id]

        # Check dependencies
        for dep_id in task.dependencies:
            dep_task = self.tasks.get(dep_id)
            if dep_task and dep_task.status != TaskStatus.COMPLETED:
                task.status = TaskStatus.BLOCKED
                return AgentResult(
                    agent_role=task.assigned_to or AgentRole.IMPLEMENTER,
                    task_id=task_id,
                    success=False,
                    output="",
                    error=f"Blocked by dependency: {dep_id}",
                )

        task.status = TaskStatus.IN_PROGRESS
        start_time = datetime.now()

        try:
            # Get agent configuration
            role = task.assigned_to or AgentRole.IMPLEMENTER
            config = AGENT_CONFIGS.get(role)

            if not config:
                raise ValueError(f"Unknown agent role: {role}")

            # Execute with model
            result = await self._run_agent(config, task.description)

            duration = (datetime.now() - start_time).total_seconds()

            agent_result = AgentResult(
                agent_role=role,
                task_id=task_id,
                success=True,
                output=result,
                duration_seconds=duration,
            )

            task.status = TaskStatus.COMPLETED
            task.completed_at = datetime.now()
            task.result = agent_result

        except Exception as e:
            duration = (datetime.now() - start_time).total_seconds()

            agent_result = AgentResult(
                agent_role=task.assigned_to or AgentRole.IMPLEMENTER,
                task_id=task_id,
                success=False,
                output="",
                error=str(e),
                duration_seconds=duration,
            )

            task.status = TaskStatus.FAILED
            task.result = agent_result

        self.results.append(agent_result)
        return agent_result

    async def _run_agent(self, config: AgentConfig, task_description: str) -> str:
        """Run an agent with the given configuration.

        Args:
            config: Agent configuration.
            task_description: Task to execute.

        Returns:
            Agent output string.
        """
        if not self.model:
            return f"[{config.name}] Would execute: {task_description}"

        # Build prompt with context
        context_str = ""
        if self.shared_context:
            context_str = "\n\nContext from previous agents:\n"
            for key, value in self.shared_context.items():
                context_str += f"- {key}: {value}\n"

        prompt = f"""{config.system_prompt}

{context_str}

Task: {task_description}

Please complete this task according to your role and responsibilities."""

        response = await asyncio.to_thread(self.model.invoke, prompt)

        output = response.content if hasattr(response, "content") else str(response)

        # Update shared context
        self.shared_context[f"{config.name}_output"] = output[:500]

        return output

    async def execute_workflow(
        self,
        task_description: str,
        workflow_type: str = "feature",
    ) -> WorkflowResult:
        """Execute a complete workflow with multiple agents.

        Args:
            task_description: Main task description.
            workflow_type: Type of workflow (feature, bugfix, etc.)

        Returns:
            WorkflowResult with all agent outputs.
        """
        start_time = datetime.now()
        workflow_result = WorkflowResult(success=True)

        # Get agents for this workflow
        agent_roles = get_agents_for_task(workflow_type)

        # Create tasks for each agent
        task_ids = []
        prev_task_id = None

        for role in agent_roles:
            config = AGENT_CONFIGS.get(role)
            if not config:
                continue

            # Create task with dependency on previous
            dependencies = [prev_task_id] if prev_task_id else []
            task_id = self.create_task(
                description=f"[{config.name}] {task_description}",
                assigned_to=role,
                dependencies=dependencies,
            )
            task_ids.append(task_id)
            prev_task_id = task_id

        # Execute tasks in order
        for task_id in task_ids:
            result = await self.execute_task(task_id)
            workflow_result.results.append(result)

            if not result.success:
                workflow_result.success = False
                workflow_result.errors.append(result.error or "Unknown error")

        # Compile final output
        outputs = []
        for result in workflow_result.results:
            outputs.append(f"## {result.agent_role.value.title()} Output\n{result.output}")

        workflow_result.final_output = "\n\n".join(outputs)
        workflow_result.total_duration = (datetime.now() - start_time).total_seconds()

        return workflow_result

    async def execute_parallel(
        self,
        task_descriptions: list[tuple[str, AgentRole]],
    ) -> list[AgentResult]:
        """Execute multiple tasks in parallel.

        Args:
            task_descriptions: List of (description, role) tuples.

        Returns:
            List of AgentResults.
        """
        # Create all tasks
        task_ids = []
        for description, role in task_descriptions:
            task_id = self.create_task(description, assigned_to=role)
            task_ids.append(task_id)

        # Execute in parallel
        results = await asyncio.gather(
            *[self.execute_task(task_id) for task_id in task_ids],
            return_exceptions=True,
        )

        # Filter results
        agent_results = []
        for result in results:
            if isinstance(result, AgentResult):
                agent_results.append(result)
            elif isinstance(result, Exception):
                agent_results.append(AgentResult(
                    agent_role=AgentRole.IMPLEMENTER,
                    task_id="unknown",
                    success=False,
                    output="",
                    error=str(result),
                ))

        return agent_results

    def get_status(self) -> dict:
        """Get orchestrator status.

        Returns:
            Status dictionary.
        """
        task_counts = {
            "pending": sum(1 for t in self.tasks.values() if t.status == TaskStatus.PENDING),
            "in_progress": sum(1 for t in self.tasks.values() if t.status == TaskStatus.IN_PROGRESS),
            "completed": sum(1 for t in self.tasks.values() if t.status == TaskStatus.COMPLETED),
            "failed": sum(1 for t in self.tasks.values() if t.status == TaskStatus.FAILED),
            "blocked": sum(1 for t in self.tasks.values() if t.status == TaskStatus.BLOCKED),
        }

        return {
            "total_tasks": len(self.tasks),
            "task_counts": task_counts,
            "results_count": len(self.results),
            "shared_context_keys": list(self.shared_context.keys()),
        }

    def reset(self) -> None:
        """Reset the orchestrator state."""
        self.tasks.clear()
        self.results.clear()
        self.shared_context.clear()
        self._task_counter = 0


# Convenience functions
async def run_multi_agent(
    task: str,
    workflow_type: str = "feature",
    model: Optional["BaseChatModel"] = None,
) -> dict:
    """Run a multi-agent workflow.

    Args:
        task: Task description.
        workflow_type: Type of workflow.
        model: Optional LLM model.

    Returns:
        Workflow result dictionary.
    """
    orchestrator = AgentOrchestrator(model)
    result = await orchestrator.execute_workflow(task, workflow_type)
    return result.to_dict()


async def delegate_to_agent(
    task: str,
    role: str,
    model: Optional["BaseChatModel"] = None,
) -> dict:
    """Delegate a task to a specific agent.

    Args:
        task: Task description.
        role: Agent role name.
        model: Optional LLM model.

    Returns:
        Agent result dictionary.
    """
    try:
        agent_role = AgentRole(role)
    except ValueError:
        return {"success": False, "error": f"Unknown role: {role}"}

    orchestrator = AgentOrchestrator(model)
    task_id = orchestrator.create_task(task, assigned_to=agent_role)
    result = await orchestrator.execute_task(task_id)
    return result.to_dict()
