"""
Task models for AgentChain.
"""

from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any
from enum import Enum
from datetime import datetime
import uuid


class TaskStatus(Enum):
    """Status of a task."""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    BLOCKED = "blocked"


class TaskPriority(Enum):
    """Task priority levels."""
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4


@dataclass
class Task:
    """
    A single task in the execution plan.
    """
    id: str
    title: str
    description: str = ""
    status: TaskStatus = TaskStatus.PENDING
    priority: TaskPriority = TaskPriority.MEDIUM
    
    # Files this task works with
    files_to_create: List[str] = field(default_factory=list)
    files_to_modify: List[str] = field(default_factory=list)
    
    # Dependencies
    depends_on: List[str] = field(default_factory=list)
    
    # Execution details
    assigned_agent: Optional[str] = None
    retry_count: int = 0
    max_retries: int = 3
    
    # Results
    output: Optional[str] = None
    error: Optional[str] = None
    
    # Metadata
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    completed_at: Optional[datetime] = None
    
    @classmethod
    def create(
        cls,
        title: str,
        description: str = "",
        **kwargs
    ) -> "Task":
        """Create a new task with auto-generated ID."""
        return cls(
            id=f"task_{uuid.uuid4().hex[:8]}",
            title=title,
            description=description,
            **kwargs
        )
    
    def can_execute(self, completed_tasks: List[str]) -> bool:
        """Check if this task can execute based on dependencies."""
        return all(dep in completed_tasks for dep in self.depends_on)
    
    def can_retry(self) -> bool:
        """Check if this task can be retried."""
        return self.retry_count < self.max_retries
    
    def mark_complete(self, output: Optional[str] = None) -> None:
        """Mark task as completed."""
        self.status = TaskStatus.COMPLETED
        self.output = output
        self.completed_at = datetime.now()
    
    def mark_failed(self, error: str) -> None:
        """Mark task as failed."""
        self.status = TaskStatus.FAILED
        self.error = error
        self.retry_count += 1
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "title": self.title,
            "description": self.description,
            "status": self.status.value,
            "priority": self.priority.value,
            "files_to_create": self.files_to_create,
            "files_to_modify": self.files_to_modify,
            "depends_on": self.depends_on,
            "assigned_agent": self.assigned_agent,
            "metadata": self.metadata,
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Task":
        """Create from dictionary."""
        return cls(
            id=data.get("id", f"task_{uuid.uuid4().hex[:8]}"),
            title=data["title"],
            description=data.get("description", ""),
            status=TaskStatus(data.get("status", "pending")),
            priority=TaskPriority(data.get("priority", 2)),
            files_to_create=data.get("files_to_create", []),
            files_to_modify=data.get("files_to_modify", []),
            depends_on=data.get("depends_on", []),
            assigned_agent=data.get("assigned_agent"),
            metadata=data.get("metadata", {}),
        )


@dataclass
class Plan:
    """
    Execution plan containing tasks.
    """
    name: str
    description: str = ""
    tasks: List[Task] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    
    def add_task(self, task: Task) -> None:
        """Add a task to the plan."""
        self.tasks.append(task)
    
    def get_task(self, task_id: str) -> Optional[Task]:
        """Get a task by ID."""
        for task in self.tasks:
            if task.id == task_id:
                return task
        return None
    
    def get_next_task(self) -> Optional[Task]:
        """Get the next executable task."""
        completed = [t.id for t in self.tasks if t.status == TaskStatus.COMPLETED]
        
        for task in self.tasks:
            if task.status == TaskStatus.PENDING and task.can_execute(completed):
                return task
        return None
    
    def is_complete(self) -> bool:
        """Check if all tasks are completed."""
        return all(t.status == TaskStatus.COMPLETED for t in self.tasks)
    
    def get_progress(self) -> Dict[str, Any]:
        """Get plan progress statistics."""
        total = len(self.tasks)
        completed = sum(1 for t in self.tasks if t.status == TaskStatus.COMPLETED)
        failed = sum(1 for t in self.tasks if t.status == TaskStatus.FAILED)
        pending = sum(1 for t in self.tasks if t.status == TaskStatus.PENDING)
        in_progress = sum(1 for t in self.tasks if t.status == TaskStatus.IN_PROGRESS)
        
        return {
            "total": total,
            "completed": completed,
            "failed": failed,
            "pending": pending,
            "in_progress": in_progress,
            "percent": round(completed / max(total, 1) * 100, 1),
        }
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "tasks": [t.to_dict() for t in self.tasks],
            "metadata": self.metadata,
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Plan":
        """Create from dictionary."""
        plan = cls(
            name=data.get("name", "Unnamed Plan"),
            description=data.get("description", ""),
            metadata=data.get("metadata", {}),
        )
        for task_data in data.get("tasks", []):
            plan.add_task(Task.from_dict(task_data))
        return plan
