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,
    ):
        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()

    def fetch_pending_tasks(self) -> List[Dict[str, Any]]:
        """Fetch pending tasks from Web backend."""
        if requests is None:
            return self._scan_local_task_files()
        try:
            resp = requests.get(
                f"{self.web_base_url}/api/tasks",
                params={"status": "pending", "limit": 10},
                timeout=10,
            )
            if resp.status_code == 200:
                data = resp.json()
                return data.get("tasks", [])
        except Exception as e:
            if self.on_error:
                self.on_error(f"Fetch tasks failed: {str(e)}")
        return []

    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)
            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 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_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
        self.dispatch_task(task_id)

        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
                point_result["params"] = params
                results.append(point_result)
            except Exception as e:
                # A2 fix: failed points are recorded as failed, NOT mock data
                results.append({
                    "point_index": idx,
                    "params": params,
                    "status": "FAILED",
                    "error": str(e),
                })
                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)

        # 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 start_polling(self, interval: int = 5) -> threading.Thread:
        """Start background thread to poll for and execute tasks."""
        self._running = True
        self._stop_event.clear()

        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)

        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 that uses RobustMotorCADSolver for real Motor-CAD simulation.

    A3/A4/A5 fixes:
    - Reuses RobustMotorCADSolver (open_new_instance=True, set_visible,
      baseline reload per point, popup suppression, write-back verification)
    - Write-back verification failures propagate (no silent except:pass)
    - Results extracted via export file parsing (not bogus get_variable names)
    - Failed points raise exception -> recorded as status=failed (no mock fallback)
    """

    def __init__(self, *args, model_path: Optional[str] = None, **kwargs):
        # Mock fallback is disabled by default for real Motor-CAD executor
        kwargs.setdefault("enable_mock", False)
        super().__init__(*args, **kwargs)
        self.model_path = model_path
        self._solver = None

    def _ensure_solver(self):
        """Lazily create RobustMotorCADSolver instance."""
        if self._solver is not None:
            return self._solver
        from robust_motorcad import RobustMotorCADSolver
        if not self.model_path:
            raise RuntimeError("model_path is required for MotorCADTaskExecutor")
        # Output directory under task dir
        output_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            "output", f"task_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        )
        self._solver = RobustMotorCADSolver(
            model_path=self.model_path,
            output_dir=output_dir,
        )
        self._solver.connect()
        return self._solver

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

        A2 fix: No mock fallback on failure. Exception propagates to
        execute_task which records status=failed.
        A3 fix: RobustMotorCADSolver handles open_new_instance, set_visible,
        baseline reload, popup suppression.
        A4 fix: Write-back verification inside solver raises on mismatch.
        A5 fix: Results from export file parsing, not get_variable.
        """
        solver = self._ensure_solver()
        # run_single_point handles baseline reload, write-verify, calculation,
        # export, parsing, and per-point disk flush.
        point_result = solver.run_single_point(params, point_index=index)
        return point_result

    def cleanup(self):
        """Disconnect solver and release Motor-CAD instance."""
        if self._solver is not None:
            try:
                self._solver.disconnect()
            except Exception:
                pass
            self._solver = 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.")