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, Literal
from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from ..database import get_db
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: Optional[Dict[str, Any]] = Field(default=None, description="Full plan data (auto-loaded from plan_id when omitted)")
    parameters: Optional[List[Dict[str, Any]]] = Field(default=None, description="Parameter sets (auto-expanded from plan_id when omitted)")
    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)")
    thermal_mode: Literal["off", "steady", "coupled"] = Field(
        default="steady",
        description="Thermal mode: off (EM only) / steady (EM + steady-state "
                    "thermal, default) / coupled (magnetic-thermal coupled).",
    )


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, db: Session = Depends(get_db)):
    """Create a new simulation task.

    When only plan_id is provided, plan_data is loaded from the plan and
    parameters are expanded from its variables (same logic as the plan's
    one-click start), so the UI no longer requires hand-written JSON.
    Explicit plan_data / parameters still take precedence (advanced override).
    """
    plan_data = request.plan_data
    parameters = request.parameters

    if (not plan_data or not parameters) and request.plan_id is not None:
        from ..models.simulation_plan import SimulationPlan
        from .plans import _expand_plan_to_parameters
        plan = db.query(SimulationPlan).filter(SimulationPlan.id == request.plan_id).first()
        if not plan:
            raise HTTPException(status_code=404, detail=f"Plan {request.plan_id} not found")
        if not plan_data:
            plan_data = plan.get_plan_dict()
        if not parameters:
            parameters = _expand_plan_to_parameters(plan_data)

    if not plan_data:
        raise HTTPException(status_code=400, detail="plan_data is required (or provide a valid plan_id)")
    if not parameters:
        raise HTTPException(status_code=400, detail="parameters is empty (plan has no valid variables to simulate)")

    try:
        manager = get_task_manager()
        task = manager.create_task(
            plan_id=request.plan_id,
            plan_data=plan_data,
            parameters=parameters,
            task_name=request.task_name,
            priority=request.priority,
            thermal_mode=request.thermal_mode,
        )
        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.delete("/{task_id}", status_code=204)
def delete_task(task_id: str):
    """Delete a finished task (terminal status only)."""
    try:
        manager = get_task_manager()
        manager.delete_task(task_id)
        return None
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


class BatchDeleteTasksRequest(BaseModel):
    """Request to delete multiple finished tasks."""
    task_ids: List[str] = Field(..., description="Task IDs to delete")


@router.post("/batch-delete")
def batch_delete_tasks(request: BatchDeleteTasksRequest):
    """Delete multiple finished tasks; per-item failures do not abort the batch."""
    manager = get_task_manager()
    deleted, errors = [], []
    for task_id in request.task_ids:
        try:
            manager.delete_task(task_id)
            deleted.append(task_id)
        except ValueError as e:
            errors.append({"task_id": task_id, "error": str(e)})
    return {"deleted": deleted, "deleted_count": len(deleted), "errors": errors}


@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")