wordcloud/backend/service/runner.py

from __future__ import annotations

import json
import logging
import os
import subprocess
import sys
import time
from pathlib import Path

from .job_manager import JobManager
from .log_config import get_logger
from .schemas import JobPaths

log = get_logger("service.runner")


STAGE_RULES: list[tuple[str, str, int]] = [
    ("--- 1. 读取数据 ---", "reading_data", 10),
    ("--- 2. 智能画幅计算 ---", "sizing_canvas", 20),
    ("--- 3. 生成掩膜", "building_mask", 35),
    ("--- 4. 计算权重 ---", "computing_weights", 45),
    ("--- 5. 启动生成", "placing_words", 65),
    ("--- 6. 高清渲染 ---", "rendering", 85),
    ("已保存:", "writing_outputs", 92),
    ("✅ 完成", "writing_outputs", 99),
]


class JobRunner:
    def __init__(self, project_root: Path, manager: JobManager) -> None:
        self.project_root = project_root
        self.manager = manager
        self.script_path = self.project_root / "wordcloud_generate_hybrid.py"

    def _parse_stage(self, line: str, current_stage: str, current_progress: int) -> tuple[str, int]:
        for token, stage, progress in STAGE_RULES:
            if token in line:
                return stage, progress

        if "尝试 #" in line or "尺度" in line or "二分重试" in line:
            progress = max(current_progress, 70)
            return "placing_words", min(progress + 1, 84)

        return current_stage, current_progress

    def run(self, job_id: str, paths: JobPaths, config: dict) -> None:
        t_start = time.time()
        log.info("=" * 50)
        log.info("[Runner] 任务启动 job_id=%s", job_id)
        log.info("  config_path = %s", paths.config_path)
        log.info("  output_dir  = %s", paths.output_dir)
        log.info("  子进程 python = %s", sys.executable)
        self.manager.set_status(job_id, status="running", stage="starting", progress_percent=1, message="任务启动")

        with paths.config_path.open("w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
        log.info("  配置文件已写入")

        cmd = [
            sys.executable,
            str(self.script_path),
            "--config",
            str(paths.config_path),
        ]

        env = os.environ.copy()
        process = subprocess.Popen(
            cmd,
            cwd=str(self.project_root),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            env=env,
        )

        stage = "starting"
        progress = 1

        assert process.stdout is not None
        for raw in process.stdout:
            line = raw.rstrip("\n")
            log.info("[Pipeline] %s", line)
            stage, progress = self._parse_stage(line, stage, progress)
            self.manager.add_event(
                job_id,
                kind="log",
                stage=stage,
                progress_percent=progress,
                message=line,
            )

        ret = process.wait()
        elapsed = time.time() - t_start
        log.info("[Runner] 子进程退出 code=%d 耗时=%.2fs", ret, elapsed)

        png = next(paths.output_dir.glob("*.png"), None)
        svg = next(paths.output_dir.glob("*[!_stroke].svg"), None)
        svg_stroke = next(paths.output_dir.glob("*_stroke.svg"), None)
        db = next(paths.output_dir.glob("*.db"), None)
        metrics = next(paths.output_dir.glob("*metrics*.json"), None)

        log.info("[Runner] 产物扫描:")
        log.info("  png        = %s", png)
        log.info("  svg        = %s", svg)
        log.info("  svg_stroke = %s", svg_stroke)
        log.info("  db         = %s", db)
        log.info("  metrics    = %s", metrics)

        artifacts = {
            "png": str(png) if png else "",
            "svg": str(svg) if svg else "",
            "svg_stroke": str(svg_stroke) if svg_stroke else "",
            "db": str(db) if db else "",
            "metrics": str(metrics) if metrics else "",
        }
        self.manager.set_artifacts(job_id, artifacts)

        if ret == 0 and not png:
            log.error("[Runner] 任务失败：退出码=0 但未找到输出图片")
            self.manager.add_event(
                job_id,
                kind="status",
                stage="failed",
                progress_percent=100,
                message="任务失败：未找到输出图片",
            )
            self.manager.set_status(
                job_id,
                status="failed",
                stage="failed",
                progress_percent=100,
                message="任务失败",
                error="missing png artifact",
            )
            return

        if ret == 0:
            log.info("[Runner] ✅ 任务完成 job_id=%s 总耗时=%.2fs", job_id, elapsed)
            self.manager.add_event(
                job_id,
                kind="status",
                stage="completed",
                progress_percent=100,
                message="任务完成",
            )
            self.manager.set_status(job_id, status="success", stage="completed", progress_percent=100, message="任务完成")
        else:
            log.error("[Runner] ❌ 任务失败 job_id=%s exit_code=%d 耗时=%.2fs", job_id, ret, elapsed)
            self.manager.add_event(
                job_id,
                kind="status",
                stage="failed",
                progress_percent=100,
                message=f"任务失败，退出码: {ret}",
            )
            self.manager.set_status(
                job_id,
                status="failed",
                stage="failed",
                progress_percent=100,
                message="任务失败",
                error=f"script exited with code {ret}",
            )