"""Task coordinator for multi-agent workflows.

This module provides task decomposition, dependency management,
and execution planning for multi-agent collaboration.
"""

from __future__ import annotations

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

from ieum.agents.roles import AgentRole, get_agents_for_task

if TYPE_CHECKING:
    from langchain_core.language_models import BaseChatModel


class TaskPriority(Enum):
    """Task priority levels."""

    CRITICAL = 0
    HIGH = 1
    MEDIUM = 2
    LOW = 3


@dataclass
class Task:
    """Represents a task in the execution plan."""

    id: str
    description: str
    agent_role: AgentRole
    priority: TaskPriority = TaskPriority.MEDIUM
    dependencies: list[str] = field(default_factory=list)
    estimated_complexity: str = "medium"  # low, medium, high
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "description": self.description,
            "agent_role": self.agent_role.value,
            "priority": self.priority.value,
            "dependencies": self.dependencies,
            "estimated_complexity": self.estimated_complexity,
            "metadata": self.metadata,
        }


@dataclass
class ExecutionPlan:
    """Execution plan for a set of tasks."""

    tasks: list[Task] = field(default_factory=list)
    execution_order: list[str] = field(default_factory=list)
    parallel_groups: list[list[str]] = field(default_factory=list)
    estimated_total_complexity: str = "medium"
    warnings: list[str] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.now)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "tasks": [t.to_dict() for t in self.tasks],
            "execution_order": self.execution_order,
            "parallel_groups": self.parallel_groups,
            "estimated_total_complexity": self.estimated_total_complexity,
            "warnings": self.warnings,
            "created_at": self.created_at.isoformat(),
        }


class TaskCoordinator:
    """Coordinates task decomposition and planning."""

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

        Args:
            model: Optional LLM for task analysis.
        """
        self.model = model
        self._task_counter = 0

    def analyze_task(self, description: str) -> dict:
        """Analyze a task to determine its characteristics.

        Args:
            description: Task description.

        Returns:
            Analysis dictionary with task_type, complexity, etc.
        """
        description_lower = description.lower()

        # Determine task type
        task_type = "feature"  # default

        if any(word in description_lower for word in ["fix", "bug", "error", "issue"]):
            task_type = "bugfix"
        elif any(word in description_lower for word in ["refactor", "restructure", "reorganize"]):
            task_type = "refactor"
        elif any(word in description_lower for word in ["document", "readme", "docs"]):
            task_type = "documentation"
        elif any(word in description_lower for word in ["test", "coverage", "unittest"]):
            task_type = "test"
        elif any(word in description_lower for word in ["review", "check", "analyze"]):
            task_type = "review"
        elif any(word in description_lower for word in ["research", "investigate", "explore"]):
            task_type = "research"

        # Estimate complexity
        complexity = "medium"
        if any(word in description_lower for word in ["simple", "small", "minor", "quick"]):
            complexity = "low"
        elif any(word in description_lower for word in ["complex", "large", "major", "extensive"]):
            complexity = "high"

        # Get recommended agents
        agents = get_agents_for_task(task_type)

        return {
            "task_type": task_type,
            "complexity": complexity,
            "recommended_agents": [a.value for a in agents],
            "requires_review": task_type in ("feature", "bugfix", "refactor"),
            "requires_testing": task_type in ("feature", "bugfix"),
        }

    def decompose_task(
        self,
        description: str,
        task_type: Optional[str] = None,
    ) -> list[Task]:
        """Decompose a task into subtasks.

        Args:
            description: Main task description.
            task_type: Optional task type override.

        Returns:
            List of subtasks.
        """
        # Analyze if type not provided
        if not task_type:
            analysis = self.analyze_task(description)
            task_type = analysis["task_type"]

        # Get agent sequence for this task type
        agents = get_agents_for_task(task_type)

        subtasks = []
        prev_task_id = None

        for agent_role in agents:
            self._task_counter += 1
            task_id = f"subtask_{self._task_counter}"

            # Create role-specific subtask description
            role_description = self._create_role_description(
                description, agent_role, task_type
            )

            subtask = Task(
                id=task_id,
                description=role_description,
                agent_role=agent_role,
                dependencies=[prev_task_id] if prev_task_id else [],
                priority=self._get_priority(agent_role, task_type),
            )

            subtasks.append(subtask)
            prev_task_id = task_id

        return subtasks

    def _create_role_description(
        self,
        main_description: str,
        role: AgentRole,
        task_type: str,
    ) -> str:
        """Create role-specific task description."""
        role_prefixes = {
            AgentRole.ARCHITECT: f"Design the architecture and approach for: {main_description}",
            AgentRole.IMPLEMENTER: f"Implement the solution for: {main_description}",
            AgentRole.REVIEWER: f"Review the implementation of: {main_description}",
            AgentRole.TESTER: f"Write tests for: {main_description}",
            AgentRole.DOCUMENTER: f"Document: {main_description}",
            AgentRole.DEBUGGER: f"Debug and fix: {main_description}",
            AgentRole.RESEARCHER: f"Research solutions for: {main_description}",
            AgentRole.PLANNER: f"Create a plan for: {main_description}",
        }

        return role_prefixes.get(role, main_description)

    def _get_priority(self, role: AgentRole, task_type: str) -> TaskPriority:
        """Determine task priority based on role and type."""
        # Critical priorities for security/bug tasks
        if task_type == "bugfix" and role in (AgentRole.DEBUGGER, AgentRole.IMPLEMENTER):
            return TaskPriority.CRITICAL

        # High priority for implementation
        if role == AgentRole.IMPLEMENTER:
            return TaskPriority.HIGH

        # Medium for review/test
        if role in (AgentRole.REVIEWER, AgentRole.TESTER):
            return TaskPriority.MEDIUM

        return TaskPriority.LOW

    def create_execution_plan(
        self,
        tasks: list[Task] | str,
    ) -> ExecutionPlan:
        """Create an execution plan for tasks.

        Args:
            tasks: List of tasks or task description string.

        Returns:
            ExecutionPlan with ordering and parallelization.
        """
        # Handle string input - decompose first
        if isinstance(tasks, str):
            tasks = self.decompose_task(tasks)

        plan = ExecutionPlan(tasks=tasks)

        # Build dependency graph
        task_map = {t.id: t for t in tasks}
        remaining = set(t.id for t in tasks)
        completed = set()

        # Determine execution order (topological sort)
        while remaining:
            # Find tasks with satisfied dependencies
            ready = []
            for task_id in remaining:
                task = task_map[task_id]
                if all(dep in completed for dep in task.dependencies):
                    ready.append(task_id)

            if not ready:
                # Circular dependency or error
                plan.warnings.append(f"Circular dependency detected among: {remaining}")
                break

            # Group ready tasks by priority for parallel execution
            parallel_group = sorted(
                ready,
                key=lambda tid: task_map[tid].priority.value
            )
            plan.parallel_groups.append(parallel_group)
            plan.execution_order.extend(parallel_group)

            # Mark as completed
            for task_id in ready:
                completed.add(task_id)
                remaining.remove(task_id)

        # Estimate total complexity
        complexities = [t.estimated_complexity for t in tasks]
        if "high" in complexities:
            plan.estimated_total_complexity = "high"
        elif all(c == "low" for c in complexities):
            plan.estimated_total_complexity = "low"
        else:
            plan.estimated_total_complexity = "medium"

        return plan

    def optimize_plan(self, plan: ExecutionPlan) -> ExecutionPlan:
        """Optimize an execution plan.

        Args:
            plan: Plan to optimize.

        Returns:
            Optimized plan.
        """
        # Identify tasks that can run in parallel
        # (same priority level, no mutual dependencies)
        task_map = {t.id: t for t in plan.tasks}

        optimized_groups = []
        for group in plan.parallel_groups:
            # Split into truly parallel tasks
            parallel_subgroups = []
            current_group = []

            for task_id in group:
                task = task_map[task_id]

                # Check if can run with current group
                can_parallel = True
                for existing_id in current_group:
                    existing = task_map[existing_id]
                    # Can't parallel if one depends on the other
                    if task_id in existing.dependencies or existing_id in task.dependencies:
                        can_parallel = False
                        break

                if can_parallel:
                    current_group.append(task_id)
                else:
                    if current_group:
                        parallel_subgroups.append(current_group)
                    current_group = [task_id]

            if current_group:
                parallel_subgroups.append(current_group)

            optimized_groups.extend(parallel_subgroups)

        plan.parallel_groups = optimized_groups
        return plan

    def estimate_duration(self, plan: ExecutionPlan) -> dict:
        """Estimate execution duration.

        Args:
            plan: Execution plan.

        Returns:
            Duration estimates.
        """
        # Rough estimates per complexity
        complexity_minutes = {
            "low": 2,
            "medium": 5,
            "high": 15,
        }

        total_sequential = 0
        total_parallel = 0

        for group in plan.parallel_groups:
            group_times = []
            for task_id in group:
                task = next(t for t in plan.tasks if t.id == task_id)
                task_time = complexity_minutes.get(task.estimated_complexity, 5)
                group_times.append(task_time)
                total_sequential += task_time

            # Parallel time is max of group
            total_parallel += max(group_times) if group_times else 0

        return {
            "sequential_minutes": total_sequential,
            "parallel_minutes": total_parallel,
            "parallelization_savings": total_sequential - total_parallel,
        }


# Convenience functions
def create_plan(
    task_description: str,
    task_type: Optional[str] = None,
) -> dict:
    """Create an execution plan for a task.

    Args:
        task_description: Task description.
        task_type: Optional task type.

    Returns:
        Execution plan dictionary.
    """
    coordinator = TaskCoordinator()
    subtasks = coordinator.decompose_task(task_description, task_type)
    plan = coordinator.create_execution_plan(subtasks)
    plan = coordinator.optimize_plan(plan)

    return {
        "plan": plan.to_dict(),
        "duration_estimate": coordinator.estimate_duration(plan),
    }


def analyze_task(task_description: str) -> dict:
    """Analyze a task.

    Args:
        task_description: Task description.

    Returns:
        Task analysis dictionary.
    """
    coordinator = TaskCoordinator()
    return coordinator.analyze_task(task_description)
