"""
BGX Group module - Task groups with dependency graph execution.

This module provides TaskGroup class that manages task dependencies
and execution patterns including sequential chains, fan-out, fan-in, and diamond patterns.
"""

import threading
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set
from enum import Enum

from .task import Task, TaskState


class GroupPattern(Enum):
    """Types of task group patterns."""
    SEQUENTIAL = "sequential"
    FAN_OUT = "fan_out"
    FAN_IN = "fan_in"
    DIAMOND = "diamond"


@dataclass
class TaskDependency:
    """Represents a dependency between tasks."""
    task_id: str
    depends_on: List[str] = field(default_factory=list)
    condition: Optional[Callable[[Any], bool]] = None


@dataclass
class GroupNode:
    """Node in the dependency graph."""
    task_id: str
    task: Optional[Task] = None
    dependencies: Set[str] = field(default_factory=set)
    dependents: Set[str] = field(default_factory=set)
    in_degree: int = 0
    out_degree: int = 0


class TaskGroup:
    """Task group with dependency graph execution support."""
    
    def __init__(self, engine: 'BackgroundEngine'):
        """Initialize task group.
        
        Args:
            engine: BackgroundEngine instance for task execution
        """
        self.engine = engine
        self._nodes: Dict[str, GroupNode] = {}
        self._edges: Dict[str, List[str]] = {}
        self._execution_queue: deque = deque()
        self._lock = threading.RLock()
        self._shutdown = False
        self._completed_tasks: Set[str] = set()
        self._failed_tasks: Set[str] = set()
        self._execution_thread = None
    
    def add_task(self, task: Task, dependencies: List[str] = None) -> str:
        """Add a task to the group.
        
        Args:
            task: Task to add
            dependencies: List of task IDs this task depends on
            
        Returns:
            Task ID for reference
        """
        with self._lock:
            if task.id in self._nodes:
                raise ValueError(f"Task {task.id} already exists in group")
            
            # Create node
            node = GroupNode(
                task_id=task.id,
                task=task,
                dependencies=set(dependencies) if dependencies else set()
            )
            
            self._nodes[task.id] = node
            self._edges[task.id] = dependencies or []
            
            # Update dependency tracking
            for dep_id in dependencies or []:
                if dep_id in self._nodes:
                    self._nodes[dep_id].dependents.add(task.id)
                    node.in_degree += 1
            
            node.out_degree = len(dependencies) if dependencies else 0
        
        return task.id
    
    def add_dependency(self, task_id: str, depends_on: str, condition: Callable[[Any], bool] = None) -> None:
        """Add a dependency between tasks.
        
        Args:
            task_id: ID of task that depends on another task
            depends_on: ID of task this task depends on
            condition: Optional condition function for conditional dependency
        """
        with self._lock:
            if task_id not in self._nodes or depends_on not in self._nodes:
                raise ValueError(f"Both tasks must exist in group")
            
            # Create dependency if not exists
            dependency = TaskDependency(task_id=task_id, depends_on=[depends_on], condition=condition)
            
            # Add to task's dependencies
            if task_id not in self._edges:
                self._edges[task_id] = []
            self._edges[task_id].append(depends_on)
            
            # Update node relationships
            self._nodes[task_id].dependencies.add(depends_on)
            self._nodes[depends_on].dependents.add(task_id)
            self._nodes[task_id].in_degree += 1
            self._nodes[depends_on].out_degree += 1
    
    def _build_execution_plan(self) -> List[str]:
        """Build execution plan using topological sort."""
        # Kahn's algorithm for topological sort
        in_degree = {task_id: node.in_degree for task_id, node in self._nodes.items()}
        queue = deque([task_id for task_id, node in self._nodes.items() if node.in_degree == 0])
        plan = []
        
        while queue:
            current = queue.popleft()
            plan.append(current)
            
            # Update in-degree of dependents
            for dependent in self._nodes[current].dependents:
                in_degree[dependent] -= 1
                if in_degree[dependent] == 0:
                    queue.append(dependent)
        
        if len(plan) != len(self._nodes):
            raise ValueError("Circular dependency detected in task group")
        
        return plan
    
    def _check_dependencies(self, task_id: str) -> bool:
        """Check if all dependencies for a task are satisfied."""
        node = self._nodes[task_id]
        
        for dep_id in node.dependencies:
            dep_node = self._nodes[dep_id]
            
            # Check if dependency is completed
            if dep_id not in self._completed_tasks:
                return False
            
            # Check if dependency failed
            if dep_id in self._failed_tasks:
                return False
            
            # Check conditional dependencies
            for edge in self._edges.get(task_id, []):
                if edge.condition and not edge.condition(dep_node.task.result):
                    return False
        
        return True
    
    def _execute_task(self, task_id: str) -> None:
        """Execute a single task."""
        node = self._nodes[task_id]
        if not node.task:
            return
        
        try:
            # Transition: PENDING -> QUEUED -> RUNNING
            if node.task.state == TaskState.PENDING:
                node.task.transition_to(TaskState.QUEUED)
            node.task.transition_to(TaskState.RUNNING)
            
            # Execute task function directly
            result = node.task.func(*node.task.args, **node.task.kwargs)
            
            # Store result back in node
            node.task.result = result
            node.task.transition_to(TaskState.SUCCESS)
            self._completed_tasks.add(task_id)
            
        except Exception as e:
            # Handle failure
            node.task.error = e
            if node.task.state == TaskState.RUNNING:
                node.task.transition_to(TaskState.FAILED)
            self._failed_tasks.add(task_id)
    
    def execute_sequential(self, tasks: List[str]) -> Dict[str, Any]:
        """Execute tasks in sequential order."""
        results = {}
        
        for task_id in tasks:
            if task_id not in self._nodes:
                raise ValueError(f"Task {task_id} not found in group")
            
            self._execute_task(task_id)
            node = self._nodes[task_id]
            results[task_id] = node.task.result
            
            # Stop execution if task failed
            if node.task.state == TaskState.FAILED:
                break
        
        return results
    
    def execute_fan_out(self, task_id: str, dependents: List[str]) -> Dict[str, Any]:
        """Execute fan-out pattern (one task triggers multiple dependents)."""
        if task_id not in self._nodes:
            raise ValueError(f"Task {task_id} not found in group")
        
        # Execute the main task
        self._execute_task(task_id)
        
        # Execute all dependents
        results = {task_id: self._nodes[task_id].task.result}
        
        for dep_id in dependents:
            if dep_id not in self._nodes:
                raise ValueError(f"Dependent task {dep_id} not found in group")
            
            self._execute_task(dep_id)
            results[dep_id] = self._nodes[dep_id].task.result
        
        return results
    
    def execute_fan_in(self, tasks: List[str]) -> Dict[str, Any]:
        """Execute fan-in pattern (multiple tasks trigger one dependent)."""
        if len(tasks) < 2:
            raise ValueError("Fan-in pattern requires at least 2 tasks")
        
        # Execute all tasks
        results = {}
        for task_id in tasks:
            if task_id not in self._nodes:
                raise ValueError(f"Task {task_id} not found in group")
            
            self._execute_task(task_id)
            results[task_id] = self._nodes[task_id].task.result
        
        return results
    
    def execute_diamond(self, start_task: str, parallel_tasks: List[str], end_task: str) -> Dict[str, Any]:
        """Execute diamond pattern (start -> parallel -> end)."""
        if start_task not in self._nodes or end_task not in self._nodes:
            raise ValueError("Start and end tasks must exist in group")
        
        if len(parallel_tasks) < 2:
            raise ValueError("Diamond pattern requires at least 2 parallel tasks")
        
        # Execute start task
        self._execute_task(start_task)
        
        # Execute parallel tasks
        results = {start_task: self._nodes[start_task].task.result}
        
        for task_id in parallel_tasks:
            if task_id not in self._nodes:
                raise ValueError(f"Parallel task {task_id} not found in group")
            
            self._execute_task(task_id)
            results[task_id] = self._nodes[task_id].task.result
        
        # Execute end task
        self._execute_task(end_task)
        results[end_task] = self._nodes[end_task].task.result
        
        return results
    
    def execute(self, pattern: GroupPattern = GroupPattern.SEQUENTIAL, task_ids: List[str] = None) -> Dict[str, Any]:
        """Execute the task group with specified pattern.
        
        Args:
            pattern: Execution pattern
            task_ids: List of task IDs to execute (None for all)
            
        Returns:
            Dictionary mapping task IDs to results
        """
        if task_ids is None:
            task_ids = list(self._nodes.keys())
        
        if not task_ids:
            return {}
        
        # Validate pattern
        if pattern == GroupPattern.SEQUENTIAL:
            return self.execute_sequential(task_ids)
        elif pattern == GroupPattern.FAN_OUT:
            if len(task_ids) < 2:
                raise ValueError("Fan-out pattern requires at least 2 tasks")
            return self.execute_fan_out(task_ids[0], task_ids[1:])
        elif pattern == GroupPattern.FAN_IN:
            return self.execute_fan_in(task_ids)
        elif pattern == GroupPattern.DIAMOND:
            if len(task_ids) < 3:
                raise ValueError("Diamond pattern requires at least 3 tasks")
            return self.execute_diamond(task_ids[0], task_ids[1:-1], task_ids[-1])
        else:
            raise ValueError(f"Unknown pattern: {pattern}")
    
    def get_task_status(self, task_id: str) -> Optional[str]:
        """Get the status of a task in the group."""
        with self._lock:
            if task_id not in self._nodes:
                return None
            
            node = self._nodes[task_id]
            if not node.task:
                return "pending"
            
            return node.task.state.value
    
    def get_group_status(self) -> Dict[str, str]:
        """Get the status of all tasks in the group."""
        with self._lock:
            status = {}
            for task_id, node in self._nodes.items():
                if not node.task:
                    status[task_id] = "pending"
                else:
                    status[task_id] = node.task.state.value
            return status
    
    def clear_completed(self) -> None:
        """Clear completed tasks tracking."""
        with self._lock:
            self._completed_tasks.clear()
            self._failed_tasks.clear()
    
    def reset(self) -> None:
        """Reset the task group."""
        with self._lock:
            self._nodes.clear()
            self._edges.clear()
            self._execution_queue.clear()
            self._completed_tasks.clear()
            self._failed_tasks.clear()
