pcb-afm-simulation-system/scripts/task_executor.py

"""Task executor for local simulation system (P4-M2).

Listens for tasks dispatched from Web backend, executes Motor-CAD
simulations via RobustMotorCADSolver, reports progress and results.

NOTE: All strings must be ASCII only. Chinese text uses \\uXXXX escapes.
"""
import json
import os
import sys
import time
import uuid
import threading
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any, Callable

try:
    import requests
except ImportError:
    requests = None

# Add scripts directory to path for robust_motorcad import
_SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
if _SCRIPTS_DIR not in sys.path:
    sys.path.insert(0, _SCRIPTS_DIR)


class TaskExecutor:
    """Executes simulation tasks dispatched from Web backend."""

    def __init__(
        self,
        web_base_url: str = "http://127.0.0.1:8000",
        task_dir: Optional[str] = None,
        on_progress: Optional[Callable] = None,
        on_complete: Optional[Callable] = None,
        on_error: Optional[Callable] = None,
        enable_mock: bool = False,
        executor_id: Optional[str] = None,
    ):
        self.web_base_url = web_base_url.rstrip("/")
        self.task_dir = task_dir or os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            "output", "tasks"
        )
        os.makedirs(self.task_dir, exist_ok=True)
        self.on_progress = on_progress
        self.on_complete = on_complete
        self.on_error = on_error
        self.enable_mock = enable_mock
        self._running = False
        self._current_task: Optional[Dict[str, Any]] = None
        self._stop_event = threading.Event()
        if executor_id is not None:
            self.executor_id = executor_id
        else:
            self.executor_id = "motorcad-executor-%s-%s" % (os.getpid(), uuid.uuid4().hex[:4])

    def fetch_pending_tasks(self) -> List[Dict[str, Any]]:
        """Fetch claimable tasks from Web backend.

        Web's start-simulation marks tasks as 'dispatched' immediately, while
        tasks created via the tasks API stay 'pending'. The executor claims
        BOTH states so every task gets picked up regardless of creation path.
        """
        if requests is None:
            return self._scan_local_task_files()
        tasks = []
        for st in ("pending", "dispatched"):
            try:
                resp = requests.get(
                    f"{self.web_base_url}/api/tasks",
                    params={"status": st, "limit": 20},
                    timeout=10,
                )
                if resp.status_code == 200:
                    data = resp.json()
                    tasks.extend(data.get("tasks", []))
            except Exception as e:
                if self.on_error:
                    self.on_error(f"Fetch tasks failed: {str(e)}")
        # De-duplicate by task_id (keep first occurrence)
        seen = set()
        result = []
        for t in tasks:
            tid = t.get("task_id")
            if tid and tid not in seen:
                seen.add(tid)
                result.append(t)
        return result

    def _hydrate_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """Fetch full task payload (parameters + plan_data) from Web backend.

        The list API only returns task metadata; the actual parameter sets
        live in the task.json file exposed by the download endpoint.
        """
        if requests is None:
            return task
        tid = task.get("task_id")
        if not tid:
            return task
        try:
            resp = requests.get(
                f"{self.web_base_url}/api/tasks/{tid}/download",
                timeout=10,
            )
            if resp.status_code == 200:
                full = resp.json()
                if isinstance(full, dict):
                    if full.get("parameters"):
                        task = {**task, **full}
                    elif task.get("_local_file"):
                        # local file fallback
                        try:
                            with open(task["_local_file"], "r", encoding="utf-8") as f:
                                local = json.load(f)
                            if local.get("parameters"):
                                task = {**task, **local}
                        except Exception:
                            pass
        except Exception as e:
            if self.on_error:
                self.on_error(f"Hydrate task {tid} failed: {str(e)}")
        return task

    def _scan_local_task_files(self) -> List[Dict[str, Any]]:
        """Scan local task directory for task files (fallback mode).

        Only picks up *_task.json files. Completed tasks are renamed
        to *_task.done.json to prevent infinite re-execution (B7 fix).
        """
        tasks = []
        for fname in os.listdir(self.task_dir):
            if fname.endswith("_task.json") and not fname.endswith("_task.done.json"):
                fpath = os.path.join(self.task_dir, fname)
                try:
                    with open(fpath, "r", encoding="utf-8") as f:
                        task = json.load(f)
                    task["_local_file"] = fpath
                    tasks.append(task)
                except Exception:
                    continue
        return tasks

    def _mark_local_task_done(self, task: Dict[str, Any]) -> None:
        """Rename completed local task file to prevent re-execution (B7 fix)."""
        fpath = task.get("_local_file")
        if fpath and os.path.exists(fpath):
            done_path = fpath.replace("_task.json", "_task.done.json")
            try:
                os.rename(fpath, done_path)
            except Exception as e:
                if self.on_error:
                    self.on_error(f"Failed to mark task done: {str(e)}")

    def dispatch_task(self, task_id: str) -> bool:
        """Mark task as dispatched on Web backend."""
        if requests is None:
            return True
        try:
            resp = requests.post(
                f"{self.web_base_url}/api/tasks/{task_id}/dispatch",
                timeout=10,
            )
            return resp.status_code in (200, 201)
        except Exception as e:
            if self.on_error:
                self.on_error(f"Dispatch task {task_id} failed: {str(e)}")
            return False

    def report_progress(
        self,
        task_id: str,
        current_point: int,
        total_points: int,
        current_params: Optional[Dict[str, Any]] = None,
        elapsed_time: Optional[float] = None,
    ) -> bool:
        """Report simulation progress to Web backend."""
        if requests is None:
            if self.on_progress:
                self.on_progress(task_id, current_point, total_points)
            return True
        try:
            payload = {
                "current_point": current_point,
                "total_points": total_points,
                "current_params": current_params,
                "elapsed_time": elapsed_time,
            }
            resp = requests.post(
                f"{self.web_base_url}/api/tasks/{task_id}/progress",
                json=payload,
                timeout=10,
            )
            return resp.status_code == 200
        except Exception as e:
            if self.on_error:
                self.on_error(f"Report progress failed: {str(e)}")
            return False

    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",
    ) -> bool:
        """Report final results to Web backend."""
        if requests is None:
            if self.on_complete:
                self.on_complete(task_id, results, metrics)
            return True
        try:
            payload = {
                "results": results,
                "metrics": metrics,
                "logs": logs,
                "duration": duration,
                "status": status,
            }
            resp = requests.post(
                f"{self.web_base_url}/api/tasks/{task_id}/results",
                json=payload,
                timeout=30,
            )
            return resp.status_code == 200
        except Exception as e:
            if self.on_error:
                self.on_error(f"Report results failed: {str(e)}")
            return False

    def _report_to_adaptive_loop(
        self, task: Dict[str, Any], results: List[Dict[str, Any]]
    ) -> None:
        """Feed an adaptive_batch task's results back into its loop.

        Maps each point result to {point_id, metrics, status} and posts to the
        loop's report-results endpoint. Best-effort: a failure here must not
        break normal task reporting (results are already stored on the task).
        """
        if requests is None:
            return
        if task.get("task_type") != "adaptive_batch":
            return
        loop_id = task.get("loop_id")
        if not loop_id:
            return
        point_results = []
        for r in results:
            pid = r.get("point_id")
            if pid is None:
                continue
            point_results.append({
                "point_id": pid,
                "metrics": r.get("metrics") or {},
                "status": "ok" if r.get("status") == "OK" else "failed",
            })
        if not point_results:
            return
        try:
            resp = requests.post(
                f"{self.web_base_url}/api/adaptive/loops/{loop_id}/report-results",
                json={"point_results": point_results},
                timeout=120,
            )
            if self.on_progress:
                self.on_progress(
                    f"Adaptive loop {loop_id}: reported {len(point_results)} "
                    f"points (HTTP {resp.status_code})"
                )
        except Exception as e:
            if self.on_error:
                self.on_error(f"Adaptive loop report failed ({loop_id}): {str(e)}")

    def execute_task(self, task: Dict[str, Any]) -> None:
        """Execute a single simulation task.

        This is a template method. Override _run_simulation_point in
        subclasses to implement actual Motor-CAD simulation.
        """
        task = self._hydrate_task(task)
        task_id = task.get("task_id", str(uuid.uuid4())[:8])
        parameters = task.get("parameters", [])
        total_points = len(parameters)
        results = []
        start_time = time.time()

        self._current_task = task
        # start-simulation already marks a task 'dispatched' at creation, while
        # tasks created via the tasks API stay 'pending'. Only claim (dispatch)
        # a task that is still pending; re-dispatching an already-dispatched
        # task is rejected by the backend (pending -> dispatched only) and
        # would otherwise be misreported as "not claimable".
        claimed = True
        if task.get("status") == "pending":
            claimed = self.dispatch_task(task_id)
        if not claimed:
            # Another instance already claimed this task; skip it so
            # parallel executors never duplicate the same simulation.
            if self.on_error:
                self.on_error("Task %s not claimable (claimed/network); skip" % task_id)
            self._current_task = None
            return

        for idx, params in enumerate(parameters):
            if self._stop_event.is_set():
                break

            elapsed = time.time() - start_time
            self.report_progress(task_id, idx, total_points, params, elapsed)

            try:
                point_result = self._run_simulation_point(params, idx)
                point_result["point_index"] = idx
                if "point_id" in params:
                    point_result["point_id"] = params["point_id"]
                point_result["params"] = params
                results.append(point_result)
            except Exception as e:
                # A2 fix: failed points are recorded as failed, NOT mock data
                failed_result = {
                    "point_index": idx,
                    "params": params,
                    "status": "FAILED",
                    "error": str(e),
                }
                if "point_id" in params:
                    failed_result["point_id"] = params["point_id"]
                results.append(failed_result)
                if self.on_error:
                    self.on_error(f"Point {idx} failed: {str(e)}")

        duration = time.time() - start_time
        metrics = self._compute_metrics(results)
        # Status reflects actual outcome: completed/cancelled/failed
        if self._stop_event.is_set():
            status = "cancelled"
        elif any(r.get("status") == "FAILED" for r in results):
            status = "completed_with_errors" if any(
                r.get("status") == "OK" for r in results
            ) else "failed"
        else:
            status = "completed"

        self.report_results(task_id, results, metrics, None, duration, status)
        self.report_progress(task_id, total_points, total_points, None, duration)

        # Adaptive-loop bridge: an adaptive_batch task belongs to a loop; feed
        # per-point results back to /adaptive/loops/{loop_id}/report-results so
        # the search advances without manual intervention (P3-M5 gap closure).
        self._report_to_adaptive_loop(task, results)

        # B7 fix: mark local task file as done to prevent re-execution
        if requests is None:
            self._mark_local_task_done(task)

        self._current_task = None
        if self.on_complete:
            self.on_complete(task_id, results, metrics)

    def _run_simulation_point(self, params: Dict[str, Any], index: int) -> Dict[str, Any]:
        """Run a single simulation point. Override in subclass.

        Mock data is ONLY returned when enable_mock=True (explicit opt-in).
        Mock results are tagged with source="mock" so they can never be
        confused with real simulation data (A2 fix).
        """
        if not self.enable_mock:
            raise RuntimeError(
                "No simulation backend configured. "
                "Use MotorCADTaskExecutor for real Motor-CAD simulation, "
                "or set enable_mock=True for testing."
            )
        import random
        rng = random.Random(index + hash(json.dumps(params, sort_keys=True)) % 10000)
        airgap = params.get("airgap_mm", 1.0)
        current = params.get("current_a", 15.0)
        return {
            "tavg_nm": round(current * 2.5 / (airgap ** 0.5) + rng.gauss(0, 0.3), 4),
            "efficiency_pct": round(88 + rng.gauss(0, 2), 2),
            "total_losses_w": round(50 + rng.gauss(0, 10), 2),
            "winding_temp_c": round(90 + rng.gauss(0, 10), 1),
            "status": "OK",
            "source": "mock",
        }

    def _compute_metrics(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Compute aggregated metrics from results."""
        ok_results = [r for r in results if r.get("status") == "OK"]
        if not ok_results:
            return {
                "total_points": len(results),
                "successful_points": 0,
                "failed_points": len(results),
            }
        metrics = {
            "total_points": len(results),
            "successful_points": len(ok_results),
            "failed_points": len(results) - len(ok_results),
        }
        for key in ["tavg_nm", "efficiency_pct", "total_losses_w", "winding_temp_c"]:
            values = [r[key] for r in ok_results if key in r]
            if values:
                metrics[f"{key}_min"] = min(values)
                metrics[f"{key}_max"] = max(values)
                metrics[f"{key}_mean"] = round(sum(values) / len(values), 4)
        return metrics

    def _send_heartbeat(self) -> None:
        """Register this executor with the Web backend (online status)."""
        if requests is None:
            return
        try:
            status = "running" if self._current_task is not None else "idle"
            current_task = None
            progress = None
            if self._current_task is not None:
                current_task = self._current_task.get("task_id")
                total = self._current_task.get("total_points") or 0
                done = self._current_task.get("completed_points") or 0
                progress = {
                    "completed_points": done,
                    "total_points": total,
                }
            requests.post(
                f"{self.web_base_url}/api/executor/heartbeat",
                json={
                    "executor_id": self.executor_id,
                    "status": status,
                    "current_task": current_task,
                    "progress": progress,
                },
                timeout=5,
            )
        except Exception:
            # Heartbeat failures are non-fatal
            pass

    def start_polling(self, interval: int = 5) -> threading.Thread:
        """Start background threads: one polls/executes tasks, one heartbeats.

        Heartbeat runs on its own thread so a long-running Motor-CAD point
        (~2 min each) never starves the heartbeat - otherwise the backend
        would mark this executor offline mid-task (observed 2026-09-04).
        """
        self._running = True
        self._stop_event.clear()

        def heartbeat_loop():
            while self._running and not self._stop_event.is_set():
                self._send_heartbeat()
                self._stop_event.wait(interval)

        def poll_loop():
            while self._running and not self._stop_event.is_set():
                try:
                    tasks = self.fetch_pending_tasks()
                    for task in tasks:
                        if self._stop_event.is_set():
                            break
                        self.execute_task(task)
                except Exception as e:
                    if self.on_error:
                        self.on_error(f"Poll loop error: {str(e)}")
                self._stop_event.wait(interval)

        hb_thread = threading.Thread(target=heartbeat_loop, daemon=True)
        hb_thread.start()
        thread = threading.Thread(target=poll_loop, daemon=True)
        thread.start()
        return thread

    def stop(self):
        """Stop the executor."""
        self._running = False
        self._stop_event.set()


class MotorCADTaskExecutor(TaskExecutor):
    """Task executor backed by a registered simulation-tool adapter.

    Uses afmcore.adapters.get_adapter(tool) so the executor never
    hard-codes a specific solver. Default tool "motorcad" wraps
    RobustMotorCADSolver (open_new_instance, set_visible, baseline
    reload per point, popup suppression, write-back verification,
    per-point disk flush).

    Result mapping: adapter returns {metrics, status, error, ...}; the
    metrics dict is flattened to the point's top level so downstream
    aggregation (TaskExecutor._compute_metrics) keeps working unchanged.
    """

    def __init__(self, *args, model_path: Optional[str] = None,
                         tool: str = "motorcad",
                         enable_thermal: bool = False, **kwargs):
        # Mock fallback is disabled by default for real solver adapter.
        kwargs.setdefault("enable_mock", False)
        super().__init__(*args, **kwargs)
        self.model_path = model_path
        self.tool = tool
        # P5-M6: pass through to the adapter so each EM point can also run a
        # steady-state thermal solve and merge thermal metrics.
        self.enable_thermal = bool(enable_thermal)
        self._adapter = None

    def _ensure_adapter(self):
        """Lazily create the tool adapter via the platform registry."""
        if self._adapter is not None:
            return self._adapter
        _root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        _src = os.path.join(_root, "src")
        if _src not in sys.path:
            sys.path.insert(0, _src)
        from afmcore.adapters import get_adapter
        # Dynamically import the adapter module matching self.tool so it
        # self-registers in ADAPTER_REGISTRY. Unknown tools rely on
        # pre-registered adapters (caller may have imported them).
        if self.tool == "motorcad":
            import afmcore.adapters.motorcad  # noqa: F401
        elif self.tool == "maxwell":
            import afmcore.adapters.maxwell  # noqa: F401
        elif self.tool == "jmag":
            import afmcore.adapters.jmag  # noqa: F401
        if not self.model_path:
            raise RuntimeError("model_path is required for MotorCADTaskExecutor")
        output_dir = os.path.join(
            _root, "output", "task_%s" % datetime.now().strftime("%Y%m%d_%H%M%S")
        )
        self._adapter = get_adapter(
            self.tool, model_path=self.model_path, output_dir=output_dir,
            enable_thermal=self.enable_thermal,
        )
        self._adapter.connect()
        return self._adapter

    def _run_simulation_point(self, params: Dict[str, Any], index: int) -> Dict[str, Any]:
        """Run one point through the adapter and flatten metrics to top level.

        The adapter owns the robust protocol (baseline reload, write-back
        verification, export parsing). A non-OK point raises so execute_task
        records status=FAILED (no mock fallback).

        When enable_mock=True the base-class mock implementation is used
        instead, so no Motor-CAD instance is launched at all (P5-M2).
        """
        if self.enable_mock:
            return super()._run_simulation_point(params, index)
        adapter = self._ensure_adapter()
        result = adapter.run_point(
            self.model_path, params=params,
            output_dir=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            tag=str(index),
        )
        if result.get("status") != "OK":
            raise RuntimeError(
                result.get("error") or ("Simulation failed (adapter status=%s)"
                % result.get("status"))
            )
        metrics = result.get("metrics") or {}
        point = dict(metrics)
        point["status"] = "OK"
        point["metrics"] = metrics
        point["error"] = result.get("error")
        point["solve_time_s"] = result.get("solve_time_s")
        return point

    def cleanup(self):
        """Disconnect the adapter and release the tool instance."""
        if self._adapter is not None:
            try:
                self._adapter.disconnect()
            except Exception:
                pass
            self._adapter = None


if __name__ == "__main__":
    # Standalone test: run executor with mock data (explicit)
    executor = TaskExecutor(
        web_base_url=os.environ.get("WEB_BASE_URL", "http://127.0.0.1:8000"),
        on_progress=lambda tid, cur, tot: print(f"[{tid}] Progress: {cur}/{tot}"),
        on_complete=lambda tid, res, met: print(f"[{tid}] Complete: {len(res)} points"),
        on_error=lambda msg: print(f"ERROR: {msg}"),
        enable_mock=True,
    )
    print("Task executor started (mock mode). Press Ctrl+C to stop.")
    try:
        thread = executor.start_polling(interval=5)
        while thread.is_alive():
            time.sleep(1)
    except KeyboardInterrupt:
        executor.stop()
        print("Executor stopped.")