getDomain/domain-api/app/services/juming_task_service.py

from __future__ import annotations

import threading
from datetime import datetime
from uuid import uuid4

from app.core.files import load_juming_records, save_juming_records
from app.services.juming_service import TaskStoppedError, crawl_juming


_JUMING_TASK_LOCK = threading.Lock()
_MAX_LOG_LINES = 400


def _now() -> str:
    return datetime.now().isoformat(sep=" ", timespec="seconds")


def list_juming_tasks() -> list[dict]:
    return load_juming_records()


def _save_tasks(tasks: list[dict]) -> None:
    save_juming_records(tasks)


def _append_log(task_id: str, message: str) -> None:
    with _JUMING_TASK_LOCK:
        tasks = load_juming_records()
        target = next((item for item in tasks if item["task_id"] == task_id), None)
        if not target:
            return
        logs = list(target.get("logs") or [])
        logs.append(f"[{_now()}] {message}")
        target["logs"] = logs[-_MAX_LOG_LINES:]
        target["updated_at"] = _now()
        _save_tasks(tasks)


def _update_task(task_id: str, **patch: object) -> dict | None:
    with _JUMING_TASK_LOCK:
        tasks = load_juming_records()
        target = next((item for item in tasks if item["task_id"] == task_id), None)
        if not target:
            return None
        target.update(patch)
        target["updated_at"] = _now()
        _save_tasks(tasks)
        return dict(target)


def _is_stop_requested(task_id: str) -> bool:
    tasks = load_juming_records()
    target = next((item for item in tasks if item["task_id"] == task_id), None)
    return bool(target and target.get("cancel_requested"))


def _set_phase(task_id: str, phase: str, message: str | None = None) -> None:
    phase_labels = {
        "queued": "排队中",
        "starting": "启动中",
        "fetching": "抓取中",
        "importing": "入库中",
        "completed": "已完成",
        "failed": "失败",
        "stopping": "停止中",
        "stopped": "已停止",
    }
    patch: dict[str, object] = {
        "phase": phase,
        "phase_label": phase_labels.get(phase, phase),
    }
    if message:
        patch["message"] = message
    _update_task(task_id, **patch)


def _log_and_track_phase(task_id: str, message: str) -> None:
    if "开始采集" in message or "正在抓取" in message:
        _set_phase(task_id, "fetching", message)
    elif "开始入库处理" in message or "入库完成" in message or "已存在域名" in message:
        _set_phase(task_id, "importing", message)
    _append_log(task_id, message)


def _run_juming_task(task_id: str, payload: dict) -> None:
    _update_task(task_id, status="running", started_at=_now(), message="聚名采集任务开始执行", cancel_requested=False)
    _set_phase(task_id, "starting", "聚名采集任务开始执行")
    _append_log(task_id, "任务已启动，正在准备读取聚名登录态")
    try:
        result = crawl_juming(
            payload,
            log=lambda message: _log_and_track_phase(task_id, str(message)),
            should_stop=lambda: _is_stop_requested(task_id),
        )
        stats = result.get("stats", {})
        _update_task(
            task_id,
            status="completed",
            completed_at=_now(),
            result=result,
            cancel_requested=False,
            message=(
                f"采集完成：抓取 {result.get('domains_found', 0)} 个域名，"
                f"新增 {stats.get('added', 0)}，已存在 {stats.get('exists', 0)}，无效 {stats.get('invalid', 0)}"
            ),
        )
        _set_phase(task_id, "completed")
        _append_log(task_id, "任务执行完成")
    except TaskStoppedError as exc:
        _update_task(
            task_id,
            status="stopped",
            completed_at=_now(),
            cancel_requested=False,
            message=str(exc),
        )
        _set_phase(task_id, "stopped")
        _append_log(task_id, "任务已按请求停止")
    except Exception as exc:
        _update_task(
            task_id,
            status="failed",
            completed_at=_now(),
            cancel_requested=False,
            message=f"采集失败：{exc}",
        )
        _set_phase(task_id, "failed")
        _append_log(task_id, f"任务执行失败：{exc}")


def create_juming_task(payload: dict) -> dict:
    task_id = uuid4().hex
    mode = str(payload.get("mode") or "delete_list").strip() or "delete_list"
    record = {
        "task_id": task_id,
        "mode": mode,
        "payload": dict(payload or {}),
        "status": "queued",
        "phase": "queued",
        "phase_label": "排队中",
        "cancel_requested": False,
        "message": "采集任务已创建，等待后台执行",
        "created_at": _now(),
        "updated_at": _now(),
        "started_at": "",
        "completed_at": "",
        "result": None,
        "logs": [f"[{_now()}] 已创建聚名采集任务，等待后台执行"],
    }

    with _JUMING_TASK_LOCK:
        tasks = load_juming_records()
        tasks.insert(0, record)
        _save_tasks(tasks)

    worker = threading.Thread(target=_run_juming_task, args=(task_id, dict(payload or {})), daemon=True)
    worker.start()
    return record


def retry_juming_task(task_id: str) -> dict:
    with _JUMING_TASK_LOCK:
        tasks = load_juming_records()
        target = next((item for item in tasks if item["task_id"] == task_id), None)
        if not target:
            raise ValueError("聚名采集任务不存在")
        if target.get("status") == "running":
            raise ValueError("聚名采集任务正在运行，不能重复执行")
        target["status"] = "queued"
        target["phase"] = "queued"
        target["phase_label"] = "排队中"
        target["cancel_requested"] = False
        target["message"] = "采集任务已重新加入队列"
        target["started_at"] = ""
        target["completed_at"] = ""
        target["updated_at"] = _now()
        target["result"] = None
        target["logs"] = [f"[{_now()}] 已重新加入队列，等待后台执行"]
        payload = dict(target.get("payload") or {})
        record = dict(target)
        _save_tasks(tasks)

    worker = threading.Thread(target=_run_juming_task, args=(task_id, payload), daemon=True)
    worker.start()
    return record


def request_stop_juming_task(task_id: str) -> dict:
    with _JUMING_TASK_LOCK:
        tasks = load_juming_records()
        target = next((item for item in tasks if item["task_id"] == task_id), None)
        if not target:
            raise ValueError("聚名采集任务不存在")
        if target.get("status") in {"completed", "failed", "stopped"}:
            raise ValueError("当前任务已结束，无需停止")
        target["cancel_requested"] = True
        target["phase"] = "stopping"
        target["phase_label"] = "停止中"
        target["message"] = "已发送停止请求，等待当前步骤安全退出"
        target["updated_at"] = _now()
        logs = list(target.get("logs") or [])
        logs.append(f"[{_now()}] 已收到停止请求，等待当前步骤安全退出")
        target["logs"] = logs[-_MAX_LOG_LINES:]
        _save_tasks(tasks)
        return dict(target)