LabelingSystem/backend/services/assignment_service.py

"""
Assignment Service for task distribution.
Handles task assignment preview and dispatch to annotators.
"""
from datetime import datetime
from typing import List, Dict, Any, Optional
from database import get_db_connection
from schemas.project import ProjectStatus
from schemas.task import (
    DispatchRequest,
    AssignmentPreviewRequest,
    AnnotatorAssignment,
    AssignmentPreviewResponse,
    DispatchResponse,
)


class AssignmentService:
    """Service for managing task assignments."""
    
    def get_annotator_workload(self, annotator_ids: Optional[List[str]] = None) -> Dict[str, Dict[str, Any]]:
        """
        Get current workload for annotators.
        
        Args:
            annotator_ids: Optional list of annotator IDs to filter
        
        Returns:
            Dictionary mapping annotator_id to workload info
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Get all annotators or filter by IDs
            if annotator_ids:
                placeholders = ",".join(["?" for _ in annotator_ids])
                cursor.execute(f"""
                    SELECT id, username, email
                    FROM users
                    WHERE role = 'annotator' AND id IN ({placeholders})
                """, annotator_ids)
            else:
                cursor.execute("""
                    SELECT id, username, email
                    FROM users
                    WHERE role = 'annotator'
                """)
            
            annotators = cursor.fetchall()
            
            workload = {}
            for annotator in annotators:
                annotator_id = annotator["id"]
                
                # Get task counts for this annotator
                cursor.execute("""
                    SELECT 
                        COUNT(*) as total_assigned,
                        SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) as completed,
                        SUM(CASE WHEN status = 'in_progress' THEN 1 ELSE 0 END) as in_progress,
                        SUM(CASE WHEN status = 'pending' THEN 1 ELSE 0 END) as pending
                    FROM tasks
                    WHERE assigned_to = ?
                """, (annotator_id,))
                
                counts = cursor.fetchone()
                
                workload[annotator_id] = {
                    "id": annotator_id,
                    "username": annotator["username"],
                    "email": annotator["email"],
                    "total_assigned": counts["total_assigned"] or 0,
                    "completed": counts["completed"] or 0,
                    "in_progress": counts["in_progress"] or 0,
                    "pending": counts["pending"] or 0,
                }
            
            return workload
    
    def preview_assignment(
        self, 
        project_id: str, 
        annotator_ids: List[str]
    ) -> AssignmentPreviewResponse:
        """
        Preview task assignment distribution.
        
        Args:
            project_id: Project ID
            annotator_ids: List of annotator IDs to assign tasks to
        
        Returns:
            Preview of how tasks would be distributed
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Check project exists and get status
            cursor.execute("""
                SELECT id, name, status FROM projects WHERE id = ?
            """, (project_id,))
            project = cursor.fetchone()
            
            if not project:
                raise ValueError(f"项目 '{project_id}' 不存在")
            
            # Get total and unassigned tasks count
            cursor.execute("""
                SELECT 
                    COUNT(*) as total,
                    SUM(CASE WHEN assigned_to IS NULL THEN 1 ELSE 0 END) as unassigned
                FROM tasks
                WHERE project_id = ?
            """, (project_id,))
            counts = cursor.fetchone()
            total_count = counts["total"] or 0
            unassigned_count = counts["unassigned"] or 0
            
            if unassigned_count == 0:
                raise ValueError("没有待分配的任务")
            
            if not annotator_ids:
                raise ValueError("请选择至少一个标注人员")
            
            # Validate annotators exist
            placeholders = ",".join(["?" for _ in annotator_ids])
            cursor.execute(f"""
                SELECT id, username FROM users
                WHERE id IN ({placeholders}) AND role = 'annotator'
            """, annotator_ids)
            valid_annotators = cursor.fetchall()
            
            if len(valid_annotators) != len(annotator_ids):
                raise ValueError("部分标注人员ID无效")
            
            # Calculate distribution (round-robin style, equal distribution)
            num_annotators = len(annotator_ids)
            base_count = unassigned_count // num_annotators
            remainder = unassigned_count % num_annotators
            
            # Get current workload
            workload = self.get_annotator_workload(annotator_ids)
            
            assignments = []
            for i, annotator in enumerate(valid_annotators):
                annotator_id = annotator["id"]
                # First 'remainder' annotators get one extra task
                task_count = base_count + (1 if i < remainder else 0)
                
                current_workload = workload.get(annotator_id, {})
                # Current workload = total assigned - completed (active tasks)
                active_tasks = current_workload.get("total_assigned", 0) - current_workload.get("completed", 0)
                
                assignments.append(AnnotatorAssignment(
                    user_id=annotator_id,
                    username=annotator["username"],
                    task_count=task_count,
                    percentage=round(task_count / unassigned_count * 100, 1) if unassigned_count > 0 else 0,
                    current_workload=active_tasks,
                ))
            
            return AssignmentPreviewResponse(
                project_id=project_id,
                total_tasks=total_count,
                unassigned_tasks=unassigned_count,
                assignments=assignments,
            )
    
    def dispatch_tasks(
        self, 
        project_id: str, 
        annotator_ids: List[str],
        admin_id: str
    ) -> DispatchResponse:
        """
        Dispatch tasks to annotators.
        
        Args:
            project_id: Project ID
            annotator_ids: List of annotator IDs to assign tasks to
            admin_id: ID of admin performing the dispatch
        
        Returns:
            Dispatch result with assignment details
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # Check project exists and status
            cursor.execute("""
                SELECT id, name, status FROM projects WHERE id = ?
            """, (project_id,))
            project = cursor.fetchone()
            
            if not project:
                raise ValueError(f"项目 '{project_id}' 不存在")
            
            current_status = ProjectStatus(project["status"]) if project["status"] else ProjectStatus.DRAFT
            
            # Only allow dispatch in ready status
            if current_status != ProjectStatus.READY:
                raise ValueError(f"只能在 ready 状态下分发任务，当前状态: {current_status.value}")
            
            # Get unassigned task IDs
            cursor.execute("""
                SELECT id FROM tasks
                WHERE project_id = ? AND assigned_to IS NULL
                ORDER BY id
            """, (project_id,))
            unassigned_tasks = [row["id"] for row in cursor.fetchall()]
            
            if not unassigned_tasks:
                raise ValueError("没有待分配的任务")
            
            if not annotator_ids:
                raise ValueError("请选择至少一个标注人员")
            
            # Validate annotators
            placeholders = ",".join(["?" for _ in annotator_ids])
            cursor.execute(f"""
                SELECT id, username FROM users
                WHERE id IN ({placeholders}) AND role = 'annotator'
            """, annotator_ids)
            valid_annotators = {row["id"]: row["username"] for row in cursor.fetchall()}
            
            if len(valid_annotators) != len(annotator_ids):
                raise ValueError("部分标注人员ID无效")
            
            # Distribute tasks
            num_annotators = len(annotator_ids)
            assignments_result = {aid: {"count": 0, "name": valid_annotators[aid]} for aid in annotator_ids}
            
            for i, task_id in enumerate(unassigned_tasks):
                annotator_id = annotator_ids[i % num_annotators]
                
                # Update task assignment
                cursor.execute("""
                    UPDATE tasks
                    SET assigned_to = ?, status = 'pending'
                    WHERE id = ?
                """, (annotator_id, task_id))
                
                assignments_result[annotator_id]["count"] += 1
            
            # Update project status to in_progress
            cursor.execute("""
                UPDATE projects
                SET status = ?, updated_at = ?
                WHERE id = ?
            """, (ProjectStatus.IN_PROGRESS.value, datetime.now(), project_id))
            
            # Build response
            total_assigned = len(unassigned_tasks)
            
            # Get updated workload
            workload = self.get_annotator_workload(annotator_ids)
            
            assignments = []
            for aid, info in assignments_result.items():
                current_workload = workload.get(aid, {})
                active_tasks = current_workload.get("total_assigned", 0) - current_workload.get("completed", 0)
                
                assignments.append(AnnotatorAssignment(
                    user_id=aid,
                    username=info["name"],
                    task_count=info["count"],
                    percentage=round(info["count"] / total_assigned * 100, 1) if total_assigned > 0 else 0,
                    current_workload=active_tasks,
                ))
            
            return DispatchResponse(
                project_id=project_id,
                success=True,
                total_assigned=total_assigned,
                assignments=assignments,
                project_status=ProjectStatus.IN_PROGRESS.value,
            )


# Singleton instance
assignment_service = AssignmentService()