pcb-afm-simulation-system/web/backend/app/services/task_manager.py

"""Task management service for Web-Local system dispatch (P4-M2).

Handles task creation, dispatch, status tracking, progress updates,
and result reception from local simulation executor.
"""
import json
import os
import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any
from pathlib import Path

from ..database import SessionLocal, Task


class TaskManager:
    """Manages simulation tasks between Web and Local systems."""

    TASK_STATUSES = [
        "pending",      # Created, waiting for dispatch
        "dispatched",   # Sent to local executor
        "running",      # Local executor is running
        "completed",    # All points completed successfully
        "failed",       # Task failed
        "cancelled",    # User cancelled
    ]

    def __init__(self, output_dir: Optional[str] = None):
        self.output_dir = output_dir or os.path.join(
            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
            "output", "tasks"
        )
        os.makedirs(self.output_dir, exist_ok=True)

    def create_task(
        self,
        plan_id: Optional[int],
        plan_data: Dict[str, Any],
        parameters: List[Dict[str, Any]],
        task_name: Optional[str] = None,
        priority: int = 5,
        created_by: str = "web",
    ) -> Dict[str, Any]:
        """Create a new simulation task.

        Args:
            plan_id: Associated plan ID (optional)
            plan_data: Full plan data (boundary conditions, topology, etc.)
            parameters: List of parameter sets to simulate
            task_name: Optional task name
            priority: Task priority (1-10, higher = more urgent)
            created_by: Creator identifier

        Returns:
            Created task dict
        """
        task_uuid = str(uuid.uuid4())[:8]
        task_name = task_name or f"task_{task_uuid}"

        task_dir = os.path.join(self.output_dir, f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{task_name}")
        os.makedirs(task_dir, exist_ok=True)

        # Write task file for local executor
        task_file = os.path.join(task_dir, "task.json")
        task_payload = {
            "task_id": task_uuid,
            "task_name": task_name,
            "plan_id": plan_id,
            "plan_data": plan_data,
            "parameters": parameters,
            "priority": priority,
            "created_at": datetime.now().isoformat(),
            "total_points": len(parameters),
        }
        with open(task_file, "w", encoding="utf-8") as f:
            json.dump(task_payload, f, ensure_ascii=False, indent=2)

        # Save to database
        with SessionLocal() as db:
            db_task = Task(
                task_id=task_uuid,
                task_name=task_name,
                plan_id=plan_id,
                status="pending",
                priority=priority,
                total_points=len(parameters),
                completed_points=0,
                task_dir=task_dir,
                task_file=task_file,
                created_by=created_by,
                created_at=datetime.now(),
            )
            db.add(db_task)
            db.commit()
            db.refresh(db_task)

        return self._task_to_dict(db_task)

    def dispatch_task(self, task_id: str) -> Dict[str, Any]:
        """Mark task as dispatched and ready for local executor.

        Args:
            task_id: Task UUID

        Returns:
            Updated task dict
        """
        with SessionLocal() as db:
            task = db.query(Task).filter(Task.task_id == task_id).first()
            if not task:
                raise ValueError(f"Task {task_id} not found")
            if task.status != "pending":
                raise ValueError(f"Task {task_id} is not pending (status: {task.status})")

            task.status = "dispatched"
            task.dispatched_at = datetime.now()
            db.commit()
            db.refresh(task)
        return self._task_to_dict(task)

    def update_progress(
        self,
        task_id: str,
        current_point: int,
        total_points: Optional[int] = None,
        current_params: Optional[Dict[str, Any]] = None,
        elapsed_time: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Update task progress from local executor.

        Args:
            task_id: Task UUID
            current_point: Current point index (0-based)
            total_points: Total points (optional, will use stored value)
            current_params: Current parameter values being simulated
            elapsed_time: Elapsed time in seconds

        Returns:
            Updated task dict
        """
        with SessionLocal() as db:
            task = db.query(Task).filter(Task.task_id == task_id).first()
            if not task:
                raise ValueError(f"Task {task_id} not found")

            if task.status in ("dispatched", "running"):
                task.status = "running"
                task.started_at = task.started_at or datetime.now()

            task.completed_points = current_point
            if total_points:
                task.total_points = total_points

            # Update progress metadata
            progress_data = {
                "current_point": current_point,
                "current_params": current_params,
                "elapsed_time": elapsed_time,
                "updated_at": datetime.now().isoformat(),
            }
            existing_progress = json.loads(task.progress_data or "{}")
            existing_progress.update(progress_data)
            task.progress_data = json.dumps(existing_progress, ensure_ascii=False)

            db.commit()
            db.refresh(task)
        return self._task_to_dict(task)

    def report_results(
        self,
        task_id: str,
        results: List[Dict[str, Any]],
        metrics: Optional[Dict[str, Any]] = None,
        logs: Optional[str] = None,
        duration: Optional[float] = None,
        status: str = "completed",
    ) -> Dict[str, Any]:
        """Report final results from local executor.

        Args:
            task_id: Task UUID
            results: List of simulation result dicts
            metrics: Aggregated metrics
            logs: Program logs
            duration: Total duration in seconds
            status: Final status (completed/failed)

        Returns:
            Updated task dict
        """
        with SessionLocal() as db:
            task = db.query(Task).filter(Task.task_id == task_id).first()
            if not task:
                raise ValueError(f"Task {task_id} not found")

            task.status = status
            task.completed_at = datetime.now()
            task.completed_points = len(results)
            if duration:
                task.duration = duration

            # Save results to file
            results_file = os.path.join(task.task_dir, "results.json")
            with open(results_file, "w", encoding="utf-8") as f:
                json.dump({"results": results, "metrics": metrics, "logs": logs}, f, ensure_ascii=False, indent=2)

            task.results_file = results_file
            if metrics:
                task.result_metrics = json.dumps(metrics, ensure_ascii=False)

            db.commit()
            db.refresh(task)
        return self._task_to_dict(task)

    def get_task(self, task_id: str) -> Optional[Dict[str, Any]]:
        """Get task by ID."""
        with SessionLocal() as db:
            task = db.query(Task).filter(Task.task_id == task_id).first()
            if not task:
                return None
        return self._task_to_dict(task)

    def list_tasks(
        self,
        status: Optional[str] = None,
        plan_id: Optional[int] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> Dict[str, Any]:
        """List tasks with filters."""
        with SessionLocal() as db:
            query = db.query(Task)
            if status:
                query = query.filter(Task.status == status)
            if plan_id:
                query = query.filter(Task.plan_id == plan_id)

            total = query.count()
            tasks = query.order_by(Task.created_at.desc()).offset(offset).limit(limit).all()

        return {
            "total": total,
            "limit": limit,
            "offset": offset,
            "tasks": [self._task_to_dict(t) for t in tasks],
        }

    def cancel_task(self, task_id: str) -> Dict[str, Any]:
        """Cancel a pending or running task."""
        with SessionLocal() as db:
            task = db.query(Task).filter(Task.task_id == task_id).first()
            if not task:
                raise ValueError(f"Task {task_id} not found")
            if task.status in ("completed", "failed", "cancelled"):
                raise ValueError(f"Task {task_id} already finished (status: {task.status})")

            task.status = "cancelled"
            task.completed_at = datetime.now()
            db.commit()
            db.refresh(task)
        return self._task_to_dict(task)

    def get_task_results(self, task_id: str) -> Optional[Dict[str, Any]]:
        """Get task results file content."""
        task = self.get_task(task_id)
        if not task or not task.get("results_file"):
            return None
        if not os.path.exists(task["results_file"]):
            return None
        with open(task["results_file"], "r", encoding="utf-8") as f:
            return json.load(f)

    def _task_to_dict(self, task: Task) -> Dict[str, Any]:
        """Convert Task ORM object to dict."""
        result = {
            "id": task.id,
            "task_id": task.task_id,
            "task_name": task.task_name,
            "plan_id": task.plan_id,
            "status": task.status,
            "priority": task.priority,
            "total_points": task.total_points,
            "completed_points": task.completed_points,
            "progress": round((task.completed_points / task.total_points * 100), 1) if task.total_points else 0,
            "task_dir": task.task_dir,
            "task_file": task.task_file,
            "results_file": task.results_file,
            "created_by": task.created_by,
            "created_at": task.created_at.isoformat() if task.created_at else None,
            "dispatched_at": task.dispatched_at.isoformat() if task.dispatched_at else None,
            "started_at": task.started_at.isoformat() if task.started_at else None,
            "completed_at": task.completed_at.isoformat() if task.completed_at else None,
            "duration": task.duration,
        }
        if task.progress_data:
            try:
                result["progress_data"] = json.loads(task.progress_data)
            except (json.JSONDecodeError, Exception):
                result["progress_data"] = None
        if task.result_metrics:
            try:
                result["result_metrics"] = json.loads(task.result_metrics)
            except (json.JSONDecodeError, Exception):
                result["result_metrics"] = None
        return result


# Global singleton
_task_manager: Optional[TaskManager] = None


def get_task_manager() -> TaskManager:
    """Get or create global TaskManager singleton."""
    global _task_manager
    if _task_manager is None:
        _task_manager = TaskManager()
    return _task_manager