"""
Task data models for the dashboard.
"""

from datetime import datetime
from enum import Enum
from typing import Any, Dict, Optional

from pydantic import BaseModel, Field

from phone_agent.model import ModelConfig


class TaskStatus(str, Enum):
    """Task execution status."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"


class TaskCreateRequest(BaseModel):
    """Request to create a new task."""

    device_id: str = Field(..., description="Target device ID")
    task: str = Field(..., description="Task description")
    max_steps: int = Field(100, description="Maximum execution steps")
    lang: str = Field("cn", description="Language (cn or en)")

    # Model config - use dict to avoid validation issues with ModelConfig
    base_url: str = Field(
        default="http://localhost:8000/v1", description="Model API base URL"
    )
    model_name: str = Field(default="autoglm-phone-9b", description="Model name")
    api_key: str = Field(default="EMPTY", description="API key")
    max_tokens: int = Field(default=3000, description="Max tokens per response")
    temperature: float = Field(default=0.0, description="Sampling temperature")
    top_p: float = Field(default=0.85, description="Top-p sampling parameter")
    frequency_penalty: float = Field(default=0.2, description="Frequency penalty")


class TaskSchema(BaseModel):
    """Task schema for API responses."""

    task_id: str = Field(..., description="Unique task identifier")
    device_id: str = Field(..., description="Target device ID")
    task: str = Field(..., description="Task description")
    status: TaskStatus = Field(..., description="Task status")
    current_step: int = Field(0, description="Current step number")
    max_steps: int = Field(100, description="Maximum steps")
    current_action: Optional[Dict[str, Any]] = Field(None, description="Current action")
    thinking: Optional[str] = Field(None, description="Current thinking/reasoning")
    started_at: datetime = Field(..., description="Task start time")
    updated_at: datetime = Field(..., description="Last update time")
    finished_at: Optional[datetime] = Field(None, description="Task completion time")
    error: Optional[str] = Field(None, description="Error message if failed")
    completion_message: Optional[str] = Field(
        None, description="Task completion message with details"
    )

    class Config:
        json_schema_extra = {
            "example": {
                "task_id": "task_123456",
                "device_id": "emulator-5554",
                "task": "Open WeChat",
                "status": "running",
                "current_step": 3,
                "max_steps": 100,
                "current_action": {"action": "Tap", "element": "WeChat icon"},
                "thinking": "Looking for WeChat icon on home screen",
                "started_at": "2024-01-09T10:00:00",
                "updated_at": "2024-01-09T10:00:15",
            }
        }


# For backward compatibility
TaskRequest = TaskCreateRequest