from typing import Any, Dict, Optional
from dataclasses import dataclass, field
from agilemind.execution import GenerationParams


@dataclass
class TaskAgent:
    """Agent configuration for a task."""

    # Background information of the task, will be used as system message for the agent
    background: Optional[str]
    # Model to use for this task, if not provided, the default model will be used
    model: Optional[str] = None
    # Whether to allow to use a tool for this task
    use_tool: bool = False
    # Configuration for the agent to use for this task
    config: Optional[GenerationParams] = None


@dataclass
class TaskResult:
    """Result of a task execution."""

    # Whether allowed to use a tool for the task
    use_tool: bool = False
    # Output of the task execution
    output: Optional[str] = None
    # Error message if the task failed
    error: Optional[str] = None


@dataclass
class Task:
    """Base class representing a task that needs to be performed by an agent.

    A task encapsulates the work that needs to be done and tracks the result of its execution.
    """

    # Name of the task
    name: str
    # Input data for the task
    input: str = None
    # Agent configuration for the task
    agent: TaskAgent = field(default_factory=TaskAgent)
    # Result of the task execution
    result: TaskResult = field(default_factory=TaskResult)
    # Status of the task, may be: pending, in_progress, completed, failed
    status: str = field(default="pending", init=False)
    # Whether to save the artifact generated by the task
    save_artifact: bool = False
    # Path to save the artifact generated by the task
    artifact_path: Optional[str] = None

    def __post_init__(self):
        """Validate task attributes after initialization."""
        if not self.name:
            raise ValueError("Task must have a name")
        if not self.input:
            raise ValueError("Task must have input data")

    def set_failed(self, error: Optional[str] = None) -> None:
        """Mark the task as failed with optional error information."""
        self.status = "failed"
        self.result.error = error or "<Task failed for unknown reason>"

    def is_complete(self) -> bool:
        """Check if the task is completed."""
        return self.status == "completed"

    def is_failed(self) -> bool:
        """Check if the task has failed."""
        return self.status == "failed"

    def __str__(self) -> str:
        """String representation of the task."""
        return f"Task({self.name}, status={self.status})"

    @classmethod
    def from_dict(cls, task_dict: Dict[str, Any]) -> "Task":
        """Create a Task instance from a dictionary."""
        return cls(**task_dict)
