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.

P0-5: When a task completes and has a plan_id, results are automatically
synced to the SimulationResult table so the Plan detail page can display them.
"""
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
from ..models.task import Task
from ..models.simulation_plan import SimulationPlan
from ..models.simulation_result import SimulationResult
from ..metrics_constants import METRIC_KEYS


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",
        task_type: str = "scan",
        loop_id: Optional[str] = None,
        batch_id: Optional[int] = None,
        point_ids: Optional[List] = None,
        dynamic: bool = False,
        thermal_mode: str = "steady",
    ) -> 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
            thermal_mode: Thermal simulation mode for every point of this task.
                "off" (electromagnetic only) / "steady" (EM + steady-state
                thermal, default) / "coupled" (magnetic-thermal coupled).

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

        # Validate thermal_mode (task-level switch, see plan/executor docs).
        _valid_thermal_modes = ("off", "steady", "coupled")
        if thermal_mode not in _valid_thermal_modes:
            raise ValueError(
                "thermal_mode must be one of %s, got %r"
                % (_valid_thermal_modes, thermal_mode)
            )

        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),
            "task_type": task_type,
            "loop_id": loop_id,
            "batch_id": batch_id,
            "point_ids": point_ids or [],
            "dynamic": bool(dynamic),
            "thermal_mode": thermal_mode,
        }
        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(),
                task_type=task_type,
                loop_id=loop_id,
                batch_id=batch_id,
                point_ids=json.dumps(point_ids or [], ensure_ascii=False),
                dynamic=1 if dynamic else 0,
            )
            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.

        Uses an atomic conditional UPDATE so exactly one executor instance
        can claim a pending task; a concurrent second claim gets a ValueError.

        Args:
            task_id: Task UUID

        Returns:
            Updated task dict

        Raises:
            ValueError: task not found, or already claimed / not pending.
        """
        from sqlalchemy import update
        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")
            result = db.execute(
                update(Task)
                .where(Task.task_id == task_id, Task.status == "pending")
                .values(status="dispatched", dispatched_at=datetime.now())
            )
            db.commit()
            if result.rowcount == 0:
                # another instance already claimed it, or the state moved on
                raise ValueError(
                    "Task %s is not pending (status: %s)"
                    % (task_id, task.status)
                )
            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)

        # P0-5: Auto-sync results to plan's SimulationResult table
        plan_id = getattr(task, 'plan_id', None)
        if plan_id is not None:
            self._sync_results_to_plan(plan_id, results)

        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 delete_task(self, task_id: str) -> None:
        """Delete a finished task (terminal status only) and its task file.

        Only completed/failed/cancelled tasks may be deleted; active tasks
        must be cancelled first so a running executor is never orphaned.
        """
        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 not in ("completed", "failed", "cancelled"):
                raise ValueError(
                    f"Task {task_id} is {task.status}; cancel it before deleting"
                )
            task_file = task.task_file
            db.delete(task)
            db.commit()
        # Best-effort removal of the on-disk task file after the row is gone.
        if task_file and os.path.exists(task_file):
            try:
                os.remove(task_file)
            except OSError:
                pass

    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,
            "task_type": task.task_type or "scan",
            "loop_id": task.loop_id,
            "batch_id": task.batch_id,
            "point_ids": json.loads(task.point_ids) if task.point_ids else [],
            "dynamic": bool(task.dynamic),
        }
        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

    def _sync_results_to_plan(self, plan_id: int, results: List[Dict[str, Any]]) -> None:
        """Sync task results to SimulationResult table (P0-5).

        Clears existing results for this plan and inserts new ones.
        Each result has params (scan params + fixed params) and metrics.
        """
        try:
            with SessionLocal() as db:
                plan = db.query(SimulationPlan).filter(SimulationPlan.id == plan_id).first()
                if not plan:
                    return

                # Clear existing results
                db.query(SimulationResult).filter(SimulationResult.plan_id == plan_id).delete()

                metric_keys = set(METRIC_KEYS)
                for idx, r in enumerate(results):
                    params = r.get("params", {}) or {}
                    # Metrics may be at top level or under "metrics" key
                    metrics = r.get("metrics", {}) or {}
                    if not metrics:
                        # Extract metric fields from top level
                        for k, v in r.items():
                            if k in metric_keys and isinstance(v, (int, float)):
                                metrics[k] = float(v)

                    status = r.get("status", "OK")
                    solve_time = float(r.get("solve_time_s", r.get("seconds", 0)) or 0)
                    error_msg = r.get("error", "") or ""

                    result = SimulationResult(
                        plan_id=plan_id,
                        run_index=idx + 1,
                        status=status,
                        solve_time_s=solve_time,
                        error_message=error_msg[:500] if error_msg else "",
                    )
                    result.set_params({k: float(v) for k, v in params.items() if isinstance(v, (int, float))})
                    result.set_metrics({k: float(v) for k, v in metrics.items() if isinstance(v, (int, float))})
                    db.add(result)

                # Update plan status based on results
                ok_count = sum(1 for r in results if r.get("status") == "OK")
                if ok_count == len(results) and len(results) > 0:
                    plan.status = "completed"
                elif ok_count > 0:
                    plan.status = "completed"
                elif len(results) > 0:
                    plan.status = "failed"
                db.commit()
        except Exception as e:
            # Sync failure should not break task completion
            print(f"[TaskManager] Failed to sync results to plan {plan_id}: {e}")

    # ------------------------------------------------------------------
    # Overview & Executor monitoring (P1-9)
    # ------------------------------------------------------------------

    def get_overview(self) -> Dict[str, Any]:
        """Get task overview statistics for monitoring dashboard."""
        with SessionLocal() as db:
            total = db.query(Task).count()
            pending = db.query(Task).filter(Task.status == "pending").count()
            dispatched = db.query(Task).filter(Task.status == "dispatched").count()
            running = db.query(Task).filter(Task.status == "running").count()
            completed = db.query(Task).filter(Task.status == "completed").count()
            failed = db.query(Task).filter(Task.status == "failed").count()

            # Recent tasks (last 20)
            recent = (
                db.query(Task)
                .order_by(Task.created_at.desc())
                .limit(20)
                .all()
            )
            recent_list = [self._task_to_dict(t) for t in recent]

            # Active (pending + dispatched + running)
            active = (
                db.query(Task)
                .filter(Task.status.in_(["pending", "dispatched", "running"]))
                .order_by(Task.priority.desc(), Task.created_at.asc())
                .all()
            )
            active_list = [self._task_to_dict(t) for t in active]

            return {
                "total": total,
                "pending": pending,
                "dispatched": dispatched,
                "running": running,
                "completed": completed,
                "failed": failed,
                "active_tasks": active_list,
                "recent_tasks": recent_list,
            }

    def executor_heartbeat(self, executor_id: str, status: str = "idle",
                           current_task: Optional[str] = None,
                           progress: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Record executor heartbeat and return next pending task if available.

        Args:
            executor_id: Unique executor identifier.
            status: executor status (idle / running / error).
            current_task: task_id currently being executed, if any.
            progress: current progress data.

        Returns:
            Dict with executor status and optionally a task to claim.
        """
        now = datetime.now()
        # Store heartbeat in memory (simple approach)
        if not hasattr(self, "_executors"):
            self._executors = {}
        self._executors[executor_id] = {
            "executor_id": executor_id,
            "status": status,
            "current_task": current_task,
            "last_heartbeat": now.isoformat(),
            "progress": progress,
        }
        return {"status": "ok", "executor_id": executor_id, "server_time": now.isoformat()}

    def get_executor_status(self) -> Dict[str, Any]:
        """Get all registered executors and their status."""
        if not hasattr(self, "_executors"):
            self._executors = {}
        now = datetime.now()
        executors = []
        for eid, info in self._executors.items():
            last_hb = info.get("last_heartbeat")
            online = False
            if last_hb:
                try:
                    hb_time = datetime.fromisoformat(last_hb)
                    online = (now - hb_time).total_seconds() < 60
                except Exception:
                    online = False
            executors.append({**info, "online": online})
        return {"executors": executors, "count": len(executors)}


# 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