"""Workflow Engine - A flexible async task orchestration system.

This module provides a comprehensive workflow engine for managing
complex task dependencies with support for I/O tasks, CPU-intensive tasks,
serial tasks, and parallel task execution.
"""

from __future__ import annotations

import asyncio
import time
from abc import ABC, abstractmethod
from collections import defaultdict, deque
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Sequence


class TaskStatus(Enum):
    """Task status enumeration"""

    PENDING = "pending"
    READY = "ready"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class TaskType(Enum):
    """Task type enumeration"""

    SERIAL = "serial"  # Serial task
    PARALLEL = "parallel"  # Parallel task
    ASYNC = "async"  # Async I/O task
    CPU = "cpu"  # CPU-intensive task


@dataclass
class TaskResult:
    """Task execution result"""

    task_id: str
    success: bool
    data: Any
    execution_time: float
    error: Exception | None = None


class Task(ABC):
    """Task abstract base class"""

    def __init__(self, task_id: str, task_type: TaskType, dependencies: list[str] | None = None, timeout: float = 30.0):
        self.task_id = task_id
        self.task_type = task_type
        self.dependencies = dependencies or []
        self.timeout = timeout
        self.status = TaskStatus.PENDING
        self.result: TaskResult | None = None
        self.start_time: float | None = None
        self.end_time: float | None = None

    def get_dependencies(self) -> list[str]:
        """Get list of dependent task IDs"""
        return self.dependencies.copy()

    def can_execute(self, completed_tasks: set[str]) -> bool:
        """Check if task can be executed (dependencies satisfied)"""
        return all(dep in completed_tasks for dep in self.dependencies)

    def update_status(self, status: TaskStatus):
        """Update task status"""
        self.status = status

    @abstractmethod
    async def execute(self, context: dict[str, TaskResult]) -> Any:
        """Execute task logic, must be implemented by subclasses"""
        pass

    def get_execution_time(self) -> float:
        """Get task execution time"""
        if self.start_time and self.end_time:
            return self.end_time - self.start_time
        return 0.0


class IOTask(Task):
    """I/O-intensive task"""

    def __init__(self, task_id: str, duration: float, dependencies: list[str] | None = None, timeout: float = 30.0):
        super().__init__(task_id, TaskType.ASYNC, dependencies, timeout)
        self.duration = duration

    async def execute(self, context: dict[str, TaskResult]) -> Any:
        """Simulate I/O operation"""
        print(f"[IO] Starting task {self.task_id}, estimated duration: {self.duration}s")
        return await asyncio.wait_for(self._execute_io(context), timeout=self.timeout)

    async def _execute_io(self, context: dict[str, TaskResult]) -> Any:
        """Internal I/O execution method"""
        await asyncio.sleep(self.duration)
        return f"IO task {self.task_id} completed, dependencies: {list(context.keys())}"


class CPUTask(Task):
    """CPU-intensive task"""

    def __init__(self, task_id: str, iterations: int, dependencies: list[str] | None = None, timeout: float = 30.0):
        super().__init__(task_id, TaskType.CPU, dependencies, timeout)
        self.iterations = iterations

    async def execute(self, context: dict[str, TaskResult]) -> Any:
        """CPU-intensive computation task"""
        print(f"[CPU] Starting task {self.task_id}, iterations: {self.iterations}")

        # Move CPU-intensive task to thread pool to avoid blocking event loop
        def cpu_intensive_work():
            result = 0
            for i in range(self.iterations):
                result += i * i
            return result

        # Use asyncio.wait_for with timeout to prevent infinite hangs
        # Use run_in_executor for Python 3.8 compatibility (asyncio.to_thread is Python 3.9+)
        loop = asyncio.get_event_loop()
        result = await asyncio.wait_for(loop.run_in_executor(None, cpu_intensive_work), timeout=self.timeout)
        return f"CPU task {self.task_id} completed, result: {result}"


class SerialTask(Task):
    """Serial task (stateful, must execute sequentially)"""

    def __init__(
        self, task_id: str, process_func: Callable, dependencies: list[str] | None = None, timeout: float = 30.0
    ):
        super().__init__(task_id, TaskType.SERIAL, dependencies, timeout)
        self.process_func = process_func
        self.state = {}

    async def execute(self, context: dict[str, TaskResult]) -> Any:
        """Execute serial task"""
        print(f"[Serial] Starting serial task {self.task_id}")

        # Collect results from dependent tasks
        inputs = {dep_id: context[dep_id].data for dep_id in self.dependencies}

        # Execute process function
        if asyncio.iscoroutinefunction(self.process_func):
            result = await self.process_func(inputs, self.state)
        else:
            result = self.process_func(inputs, self.state)

        # Update state
        self.state = {"last_result": result, "executed": True}

        return f"Serial task {self.task_id} completed, result: {result}"


class ParallelTask(Task):
    """Parallel task (can execute concurrently with other tasks)"""

    def __init__(
        self,
        task_id: str,
        subtasks: Sequence[Task],
        dependencies: list[str] | None = None,
        timeout: float = 30.0,
        max_concurrent: int = 3,
    ):
        super().__init__(task_id, TaskType.PARALLEL, dependencies, timeout)
        self.subtasks = subtasks
        self.max_concurrent = max_concurrent

    async def execute(self, context: dict[str, TaskResult]) -> Any:
        """Execute subtasks in parallel"""
        print(f"[Parallel] Starting parallel task {self.task_id}, contains {len(self.subtasks)} subtasks")

        # Create semaphore to control concurrency
        semaphore = asyncio.Semaphore(self.max_concurrent)

        async def execute_subtask(subtask: Task, sem: asyncio.Semaphore):
            async with sem:
                subtask.start_time = time.time()
                subtask.update_status(TaskStatus.RUNNING)

                try:
                    data = await asyncio.wait_for(subtask.execute(context), timeout=subtask.timeout)
                    subtask.result = TaskResult(
                        task_id=subtask.task_id,
                        success=True,
                        data=data,
                        execution_time=time.time() - subtask.start_time,
                    )
                    subtask.update_status(TaskStatus.COMPLETED)
                    return subtask.result
                except asyncio.TimeoutError as e:
                    error = TimeoutError(f"Task {subtask.task_id} execution timeout")
                    subtask.result = TaskResult(
                        task_id=subtask.task_id,
                        success=False,
                        data=None,
                        execution_time=time.time() - subtask.start_time,
                        error=error,
                    )
                    subtask.update_status(TaskStatus.FAILED)
                    raise error from e
                except Exception as e:
                    subtask.result = TaskResult(
                        task_id=subtask.task_id,
                        success=False,
                        data=None,
                        execution_time=time.time() - subtask.start_time,
                        error=e,
                    )
                    subtask.update_status(TaskStatus.FAILED)
                    raise e

        # Execute all subtasks in parallel
        tasks = [execute_subtask(subtask, semaphore) for subtask in self.subtasks]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Process results using zip for better readability
        successful_results = []
        failed_results = []

        for subtask, result in zip(self.subtasks, results):
            if isinstance(result, Exception):
                failed_results.append(f"Subtask {subtask.task_id} failed: {result}")
            elif isinstance(result, TaskResult):
                successful_results.append(result.data)

        if failed_results:
            return f"Parallel task {self.task_id} partially failed: {failed_results}"

        return f"Parallel task {self.task_id} completed, results: {successful_results}"


class WorkflowEngine:
    """Workflow engine - core orchestrator"""

    def __init__(self, max_concurrent: int = 4):
        self.tasks: dict[str, Task] = {}
        self.results: dict[str, TaskResult] = {}
        self.max_concurrent = max_concurrent
        self.execution_order: list[list[str]] = []

    def add_task(self, task: Task):
        """Add task to workflow"""
        self.tasks[task.task_id] = task

    def validate_dependencies(self) -> bool:
        """Validate task dependencies, ensure no circular dependencies"""
        # Build adjacency list
        graph = defaultdict(list)
        in_degree = dict.fromkeys(self.tasks, 0)

        for task_id, task in self.tasks.items():
            for dep in task.get_dependencies():
                if dep not in self.tasks:
                    raise ValueError(f"Task {task_id} depends on unknown task {dep}")
                graph[dep].append(task_id)
                in_degree[task_id] += 1

        # Detect circular dependencies
        visited = 0
        queue = deque([task_id for task_id, degree in in_degree.items() if degree == 0])

        while queue:
            current = queue.popleft()
            visited += 1

            for neighbor in graph[current]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        if visited != len(self.tasks):
            raise ValueError("Circular dependency detected in workflow")

        return True

    def calculate_execution_order(self) -> list[list[str]]:
        """Calculate task execution order (topological sort + level grouping)"""
        if not self.tasks:
            return []

        # Build adjacency list
        graph = defaultdict(list)
        in_degree = dict.fromkeys(self.tasks, 0)

        for task_id, task in self.tasks.items():
            for dep in task.get_dependencies():
                graph[dep].append(task_id)
                in_degree[task_id] += 1

        # Level-based topological sort
        execution_order = []
        queue = deque([task_id for task_id, degree in in_degree.items() if degree == 0])

        while queue:
            level_size = len(queue)
            current_level = []

            for _ in range(level_size):
                task_id = queue.popleft()
                current_level.append(task_id)

                for neighbor in graph[task_id]:
                    in_degree[neighbor] -= 1
                    if in_degree[neighbor] == 0:
                        queue.append(neighbor)

            if current_level:
                execution_order.append(current_level)

        self.execution_order = execution_order
        return execution_order

    async def execute_workflow(self) -> dict[str, TaskResult]:
        """Execute entire workflow"""
        print("=" * 50)
        print("Starting workflow execution")
        print("=" * 50)

        # Validate dependencies
        self.validate_dependencies()

        # Calculate execution order
        execution_order = self.calculate_execution_order()
        print(f"Execution plan ({len(execution_order)} phases):")
        for i, level in enumerate(execution_order, 1):
            print(f"  Phase {i}: {level}")

        # Execute by level
        completed_tasks: set[str] = set()

        for level_index, level in enumerate(execution_order, 1):
            print(f"\n{'=' * 20} Phase {level_index} ({len(level)} tasks) {'=' * 20}")

            # Filter executable tasks in this level
            ready_tasks = []
            for task_id in level:
                task = self.tasks[task_id]
                if task.can_execute(completed_tasks):
                    task.update_status(TaskStatus.READY)
                    ready_tasks.append(task)

            if not ready_tasks:
                continue

            # Create semaphore for this level to control concurrency
            semaphore = asyncio.Semaphore(self.max_concurrent)

            async def execute_single_task(task: Task, sem: asyncio.Semaphore):
                async with sem:
                    task.start_time = time.time()
                    task.update_status(TaskStatus.RUNNING)

                    try:
                        # Collect results from dependent tasks
                        dependency_results = {dep_id: self.results[dep_id] for dep_id in task.get_dependencies()}

                        # Execute task
                        data = await asyncio.wait_for(task.execute(dependency_results), timeout=task.timeout)

                        task.end_time = time.time()
                        task.result = TaskResult(
                            task_id=task.task_id, success=True, data=data, execution_time=task.get_execution_time()
                        )
                        task.update_status(TaskStatus.COMPLETED)

                        # Store result
                        self.results[task.task_id] = task.result
                        completed_tasks.add(task.task_id)

                        print(f"[OK] Task {task.task_id} completed, duration: {task.get_execution_time():.2f}s")

                        return task.result

                    except asyncio.TimeoutError as e:
                        task.end_time = time.time()
                        task.result = TaskResult(
                            task_id=task.task_id,
                            success=False,
                            data=None,
                            execution_time=task.get_execution_time(),
                            error=e,
                        )
                        task.update_status(TaskStatus.FAILED)

                        # Store result and mark as completed (even if failed)
                        self.results[task.task_id] = task.result
                        completed_tasks.add(task.task_id)

                        print(f"[FAIL] Task {task.task_id} timeout")
                        raise e from e
                    except Exception as e:
                        task.end_time = time.time()
                        task.result = TaskResult(
                            task_id=task.task_id,
                            success=False,
                            data=None,
                            execution_time=task.get_execution_time(),
                            error=e,
                        )
                        task.update_status(TaskStatus.FAILED)

                        # Store result and mark as completed (even if failed)
                        self.results[task.task_id] = task.result
                        completed_tasks.add(task.task_id)

                        print(f"[FAIL] Task {task.task_id} failed: {e}")
                        raise e

            # Execute all ready tasks in this level in parallel
            tasks_to_execute = [execute_single_task(task, semaphore) for task in ready_tasks]

            # Use return_exceptions=True to ensure all tasks complete even if some fail
            await asyncio.gather(*tasks_to_execute, return_exceptions=True)

        print(f"\n{'=' * 50}")
        print("Workflow execution completed")
        print(f"{'=' * 50}")

        return self.results

    def get_execution_summary(self) -> dict[str, Any]:
        """Get execution summary"""
        total_tasks = len(self.tasks)
        completed = sum(1 for task in self.tasks.values() if task.status == TaskStatus.COMPLETED)
        failed = sum(1 for task in self.tasks.values() if task.status == TaskStatus.FAILED)

        total_time = 0.0
        for task in self.tasks.values():
            if task.result:
                total_time += task.result.execution_time

        return {
            "total_tasks": total_tasks,
            "completed": completed,
            "failed": failed,
            "pending": total_tasks - completed - failed,
            "total_execution_time": total_time,
            "success_rate": completed / total_tasks if total_tasks > 0 else 0,
        }
