from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Literal
from datetime import datetime
import uuid


class TaskResult(BaseModel):
    """
    Represents the result of an agent's task execution.
    """

    agent: str = Field(..., description="The agent who executed the task.")
    step: int = Field(..., description="The step number associated with the task.")
    substep: Optional[float] = Field(
        None, description="The substep number (if applicable)."
    )
    result: str = Field(
        ..., description="The response or outcome of the task execution."
    )
    timestamp: datetime = Field(
        default_factory=datetime.now,
        description="Timestamp of when the result was recorded.",
    )


class SubStep(BaseModel):
    substep: float = Field(
        ..., description="Substep identifier (float, e.g., 1.1, 2.3)."
    )
    description: str = Field(..., description="Detailed action to be performed.")
    status: Literal["not_started", "in_progress", "blocked", "completed"] = Field(
        ..., description="Current state of the sub-step."
    )


class PlanStep(BaseModel):
    step: int = Field(..., description="Step identifier (integer).")
    description: str = Field(..., description="Detailed action to be performed.")
    status: Literal["not_started", "in_progress", "blocked", "completed"] = Field(
        ..., description="Current state of the step."
    )
    substeps: Optional[List[SubStep]] = Field(
        None, description="Optional list of sub-steps."
    )


class LLMWorkflowMessage(BaseModel):
    """Represents a message exchanged within the workflow."""

    id: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="Unique identifier for the message",
    )
    role: str = Field(
        ..., description="The role of the message sender, e.g., 'user' or 'assistant'"
    )
    content: str = Field(..., description="Content of the message")
    timestamp: datetime = Field(
        default_factory=datetime.now,
        description="Timestamp when the message was created",
    )
    name: Optional[str] = Field(
        default=None,
        description="Optional name of the assistant or user sending the message",
    )


class LLMWorkflowEntry(BaseModel):
    """Represents a workflow and its associated data."""

    input: str = Field(
        ..., description="The input or description of the Workflow to be performed"
    )
    output: Optional[str] = Field(
        None, description="The output or result of the Workflow, if completed"
    )
    start_time: datetime = Field(
        default_factory=datetime.now,
        description="Timestamp when the workflow was started",
    )
    end_time: Optional[datetime] = Field(
        None, description="Timestamp when the workflow was completed or failed"
    )
    messages: List[LLMWorkflowMessage] = Field(
        default_factory=list, description="Messages exchanged during the workflow"
    )
    last_message: Optional[LLMWorkflowMessage] = Field(
        default=None, description="Last processed message in the workflow"
    )
    plan: List[PlanStep] = Field(
        default_factory=list, description="Structured execution plan for the workflow."
    )
    task_history: List[TaskResult] = Field(
        default_factory=list,
        description="A history of task executions and their results.",
    )


class LLMWorkflowState(BaseModel):
    """Represents the state of multiple LLM workflows."""

    instances: Dict[str, LLMWorkflowEntry] = Field(
        default_factory=dict,
        description="Workflow entries indexed by their instance_id.",
    )
