pcb-afm-simulation-system/web/backend/app/routers/tasks.py

"""Task management router for Web-Local system dispatch (P4-M2)."""
from typing import Optional, List, Dict, Any
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from ..services.task_manager import get_task_manager

router = APIRouter(prefix="/api/tasks", tags=["Tasks"])


class CreateTaskRequest(BaseModel):
    """Request to create a new simulation task."""
    plan_id: Optional[int] = Field(default=None, description="Associated plan ID")
    plan_data: Dict[str, Any] = Field(..., description="Full plan data")
    parameters: List[Dict[str, Any]] = Field(..., description="List of parameter sets to simulate")
    task_name: Optional[str] = Field(default=None, description="Optional task name")
    priority: int = Field(default=5, ge=1, le=10, description="Task priority (1-10)")


class ProgressUpdateRequest(BaseModel):
    """Request to update task progress."""
    current_point: int = Field(..., ge=0, description="Current point index (0-based)")
    total_points: Optional[int] = Field(default=None, description="Total points")
    current_params: Optional[Dict[str, Any]] = Field(default=None)
    elapsed_time: Optional[float] = Field(default=None)


class ResultsReportRequest(BaseModel):
    """Request to report final results."""
    results: List[Dict[str, Any]] = Field(..., description="List of simulation results")
    metrics: Optional[Dict[str, Any]] = Field(default=None)
    logs: Optional[str] = Field(default=None)
    duration: Optional[float] = Field(default=None)
    status: str = Field(default="completed", description="completed / failed")


@router.post("")
def create_task(request: CreateTaskRequest):
    """Create a new simulation task."""
    try:
        manager = get_task_manager()
        task = manager.create_task(
            plan_id=request.plan_id,
            plan_data=request.plan_data,
            parameters=request.parameters,
            task_name=request.task_name,
            priority=request.priority,
        )
        return task
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Task creation failed: {str(e)}")


@router.get("")
def list_tasks(
    status: Optional[str] = None,
    plan_id: Optional[int] = None,
    limit: int = 50,
    offset: int = 0,
):
    """List tasks with optional filters."""
    manager = get_task_manager()
    return manager.list_tasks(status=status, plan_id=plan_id, limit=limit, offset=offset)


@router.get("/{task_id}")
def get_task(task_id: str):
    """Get task details by ID."""
    manager = get_task_manager()
    task = manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
    return task


@router.post("/{task_id}/dispatch")
def dispatch_task(task_id: str):
    """Mark task as dispatched (ready for local executor)."""
    try:
        manager = get_task_manager()
        return manager.dispatch_task(task_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/{task_id}/progress")
def update_progress(task_id: str, request: ProgressUpdateRequest):
    """Update task progress from local executor."""
    try:
        manager = get_task_manager()
        return manager.update_progress(
            task_id=task_id,
            current_point=request.current_point,
            total_points=request.total_points,
            current_params=request.current_params,
            elapsed_time=request.elapsed_time,
        )
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/{task_id}/results")
def report_results(task_id: str, request: ResultsReportRequest):
    """Report final results from local executor."""
    try:
        manager = get_task_manager()
        return manager.report_results(
            task_id=task_id,
            results=request.results,
            metrics=request.metrics,
            logs=request.logs,
            duration=request.duration,
            status=request.status,
        )
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/{task_id}/results")
def get_task_results(task_id: str):
    """Get task results file content."""
    manager = get_task_manager()
    results = manager.get_task_results(task_id)
    if results is None:
        raise HTTPException(status_code=404, detail=f"Results for task {task_id} not found")
    return results


@router.post("/{task_id}/cancel")
def cancel_task(task_id: str):
    """Cancel a pending or running task."""
    try:
        manager = get_task_manager()
        return manager.cancel_task(task_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/{task_id}/download")
def download_task_file(task_id: str):
    """Download task JSON file for local executor."""
    manager = get_task_manager()
    task = manager.get_task(task_id)
    if not task or not task.get("task_file"):
        raise HTTPException(status_code=404, detail=f"Task file for {task_id} not found")
    import os
    if not os.path.exists(task["task_file"]):
        raise HTTPException(status_code=404, detail="Task file not found on disk")
    from fastapi.responses import FileResponse
    return FileResponse(task["task_file"], filename=f"{task_id}_task.json", media_type="application/json")