LabelingSystem/backend/schemas/task.py

"""
Pydantic schemas for Task API.
Defines request and response models for task operations.
"""
from datetime import datetime
from typing import Optional, Any, List
from pydantic import BaseModel, Field


class TaskCreate(BaseModel):
    """Schema for creating a new task."""
    
    project_id: str = Field(..., min_length=1, description="Project ID this task belongs to")
    name: str = Field(..., min_length=1, description="Task name")
    data: dict = Field(..., description="Task data (JSON object)")
    assigned_to: Optional[str] = Field(None, description="User ID assigned to this task")
    
    class Config:
        json_schema_extra = {
            "example": {
                "project_id": "proj_123abc",
                "name": "Annotate Image Batch 1",
                "data": {
                    "image_url": "https://example.com/image1.jpg",
                    "metadata": {"batch": 1}
                },
                "assigned_to": "user_001"
            }
        }


class TaskUpdate(BaseModel):
    """Schema for updating an existing task."""
    
    name: Optional[str] = Field(None, min_length=1, description="Task name")
    data: Optional[dict] = Field(None, description="Task data (JSON object)")
    status: Optional[str] = Field(None, description="Task status (pending, in_progress, completed)")
    assigned_to: Optional[str] = Field(None, description="User ID assigned to this task")
    
    class Config:
        json_schema_extra = {
            "example": {
                "name": "Updated Task Name",
                "status": "in_progress",
                "assigned_to": "user_002"
            }
        }


class TaskResponse(BaseModel):
    """Schema for task response."""
    
    id: str = Field(..., description="Task unique identifier")
    project_id: str = Field(..., description="Project ID this task belongs to")
    name: str = Field(..., description="Task name")
    data: dict = Field(..., description="Task data (JSON object)")
    status: str = Field(..., description="Task status (pending, in_progress, completed)")
    assigned_to: Optional[str] = Field(None, description="User ID assigned to this task")
    created_at: datetime = Field(..., description="Task creation timestamp")
    progress: float = Field(default=0.0, description="Task completion progress (0.0 to 1.0)")
    
    class Config:
        json_schema_extra = {
            "example": {
                "id": "task_456def",
                "project_id": "proj_123abc",
                "name": "Annotate Image Batch 1",
                "data": {
                    "image_url": "https://example.com/image1.jpg",
                    "metadata": {"batch": 1}
                },
                "status": "in_progress",
                "assigned_to": "user_001",
                "created_at": "2024-01-12T11:00:00",
                "progress": 0.5
            }
        }


class TaskAssignRequest(BaseModel):
    """Schema for assigning a task to a user."""
    
    user_id: str = Field(..., min_length=1, description="User ID to assign the task to")
    
    class Config:
        json_schema_extra = {
            "example": {
                "user_id": "user_001"
            }
        }


class BatchAssignRequest(BaseModel):
    """Schema for batch assigning tasks to users."""
    
    task_ids: List[str] = Field(..., min_length=1, description="List of task IDs to assign")
    user_ids: List[str] = Field(..., min_length=1, description="List of user IDs to assign to")
    mode: str = Field(
        default="round_robin",
        description="Assignment mode: 'round_robin' (轮询分配) or 'equal' (平均分配)"
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "task_ids": ["task_001", "task_002", "task_003"],
                "user_ids": ["user_001", "user_002"],
                "mode": "equal"
            }
        }


class TaskAssignmentResponse(BaseModel):
    """Schema for task assignment response."""
    
    task_id: str = Field(..., description="Task ID")
    assigned_to: str = Field(..., description="User ID assigned to")
    assigned_by: str = Field(..., description="User ID who made the assignment")
    assigned_at: datetime = Field(..., description="Assignment timestamp")


class BatchAssignResponse(BaseModel):
    """Schema for batch assignment response."""
    
    success_count: int = Field(..., description="Number of successfully assigned tasks")
    failed_count: int = Field(..., description="Number of failed assignments")
    assignments: List[TaskAssignmentResponse] = Field(
        default_factory=list,
        description="List of successful assignments"
    )
    errors: List[dict] = Field(
        default_factory=list,
        description="List of errors for failed assignments"
    )


class MyTasksResponse(BaseModel):
    """Schema for current user's tasks response."""
    
    tasks: List[TaskResponse] = Field(..., description="List of tasks assigned to current user")
    total: int = Field(..., description="Total number of tasks")
    completed: int = Field(..., description="Number of completed tasks")
    in_progress: int = Field(..., description="Number of in-progress tasks")
    pending: int = Field(..., description="Number of pending tasks")


# ============== 任务分发相关Schema ==============

class AnnotatorWorkload(BaseModel):
    """标注人员工作负载"""
    user_id: str = Field(..., description="用户ID")
    username: str = Field(..., description="用户名")
    current_tasks: int = Field(..., description="当前任务数")
    completed_today: int = Field(default=0, description="今日完成数")
    
    class Config:
        json_schema_extra = {
            "example": {
                "user_id": "user_001",
                "username": "annotator1",
                "current_tasks": 15,
                "completed_today": 5
            }
        }


class AnnotatorAssignment(BaseModel):
    """单个标注人员的分配信息"""
    user_id: str = Field(..., description="用户ID")
    username: str = Field(..., description="用户名")
    task_count: int = Field(..., description="将分配的任务数")
    percentage: float = Field(..., description="分配百分比")
    current_workload: int = Field(default=0, description="当前已有任务数")
    
    class Config:
        json_schema_extra = {
            "example": {
                "user_id": "user_001",
                "username": "annotator1",
                "task_count": 50,
                "percentage": 50.0,
                "current_workload": 10
            }
        }


class AssignmentPreviewRequest(BaseModel):
    """分配预览请求"""
    user_ids: List[str] = Field(..., min_length=1, description="标注人员ID列表")
    
    class Config:
        json_schema_extra = {
            "example": {
                "user_ids": ["user_001", "user_002"]
            }
        }


class AssignmentPreviewResponse(BaseModel):
    """分配预览响应"""
    project_id: str = Field(..., description="项目ID")
    total_tasks: int = Field(..., description="总任务数")
    unassigned_tasks: int = Field(..., description="未分配任务数")
    assignments: List[AnnotatorAssignment] = Field(..., description="分配预览列表")
    
    class Config:
        json_schema_extra = {
            "example": {
                "project_id": "proj_123abc",
                "total_tasks": 100,
                "unassigned_tasks": 100,
                "assignments": [
                    {
                        "user_id": "user_001",
                        "username": "annotator1",
                        "task_count": 50,
                        "percentage": 50.0,
                        "current_workload": 10
                    }
                ]
            }
        }


class DispatchRequest(BaseModel):
    """一键分发请求"""
    user_ids: List[str] = Field(..., min_length=1, description="标注人员ID列表")
    mode: str = Field(default="equal", description="分配模式：equal（平均分配）或 round_robin（轮询分配）")
    
    class Config:
        json_schema_extra = {
            "example": {
                "user_ids": ["user_001", "user_002"],
                "mode": "equal"
            }
        }


class DispatchResponse(BaseModel):
    """分发结果响应"""
    project_id: str = Field(..., description="项目ID")
    success: bool = Field(..., description="是否成功")
    total_assigned: int = Field(..., description="分配的任务总数")
    assignments: List[AnnotatorAssignment] = Field(..., description="分配结果列表")
    project_status: str = Field(..., description="更新后的项目状态")
    
    class Config:
        json_schema_extra = {
            "example": {
                "project_id": "proj_123abc",
                "success": True,
                "total_assigned": 100,
                "assignments": [
                    {
                        "user_id": "user_001",
                        "username": "annotator1",
                        "task_count": 50,
                        "percentage": 50.0,
                        "current_workload": 60
                    }
                ],
                "project_status": "in_progress"
            }
        }