"""
Agent workflow engine data models.
"""
from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any, Awaitable, Callable, Dict, List, Optional

from pydantic import BaseModel, Field


class WorkflowTemplateKind(str, Enum):
    BUILTIN_CHAT = "builtin.chat"


class RunStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class RunEventType(str, Enum):
    STATE_ENTER = "state_enter"
    TOOL_START = "tool_start"
    TOOL_END = "tool_end"
    WORLD_BOOK_ACTIVE = "worldbook_active"
    TASKS_CREATED = "tasks_created"
    CHUNK = "chunk"
    ERROR = "error"
    COMPLETE = "complete"


class ToolSpec(BaseModel):
    name: str
    description: str = ""
    parameters: Dict[str, Any] = Field(default_factory=dict)


class SkillManifest(BaseModel):
    id: str
    name: str = ""
    description: str = ""
    path: str = ""


class WorkflowTemplate(BaseModel):
    id: str
    kind: WorkflowTemplateKind
    name: str = ""
    description: str = ""
    version: str = "1.0.0"
    state_machine_path: str = "state_machine.json"
    skills: List[SkillManifest] = Field(default_factory=list)


class ChatRunBinding(BaseModel):
    role_name: str
    chat_name: str
    template_id: str = WorkflowTemplateKind.BUILTIN_CHAT.value


class TurnCallbacks(BaseModel):
    """Optional async callbacks for streaming / WS events."""

    model_config = {"arbitrary_types_allowed": True}

    on_worldbook_active: Optional[Callable[[List[Any]], Awaitable[None]]] = None
    on_tasks_created: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None
    on_chunk: Optional[Callable[[str], Awaitable[None]]] = None


class TurnContext(BaseModel):
    """Mutable per-turn execution context passed between tools."""

    model_config = {"arbitrary_types_allowed": True}

    request_data: Dict[str, Any] = Field(default_factory=dict)
    template_id: str = WorkflowTemplateKind.BUILTIN_CHAT.value
    run_id: str = ""
    stream: bool = False
    callbacks: Optional[TurnCallbacks] = None

    current_role: str = ""
    current_chat: str = ""
    user_message: str = ""
    preset_name: Optional[str] = None
    character: Any = None
    active_entries: List[Any] = Field(default_factory=list)
    chat_history: List[Any] = Field(default_factory=list)
    prompt_messages: List[Any] = Field(default_factory=list)
    generated_content: str = ""
    token_usage: Dict[str, Any] = Field(default_factory=dict)
    duration: float = 0.0
    task_ids: Dict[str, Optional[str]] = Field(default_factory=dict)
    error: Optional[str] = None


class WorkflowRun(BaseModel):
    id: str
    template_id: str
    binding: ChatRunBinding
    status: RunStatus = RunStatus.PENDING
    started_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    finished_at: Optional[str] = None
    current_state: Optional[str] = None
    result_content: str = ""
    error: Optional[str] = None


class RunEvent(BaseModel):
    run_id: str
    type: RunEventType
    timestamp: str = Field(default_factory=lambda: datetime.now().isoformat())
    state: Optional[str] = None
    tool: Optional[str] = None
    payload: Dict[str, Any] = Field(default_factory=dict)


class ChatTurnResult(BaseModel):
    success: bool
    content: str = ""
    error: Optional[str] = None
    active_entries: List[Any] = Field(default_factory=list)
    task_ids: Dict[str, Optional[str]] = Field(default_factory=dict)
    run_id: str = ""
    workflow_template_id: str = WorkflowTemplateKind.BUILTIN_CHAT.value