Nexus/server/application/services/server_batch_service.py

"""Background server batch jobs (category, health check, Agent ops, path detect)."""

from __future__ import annotations

import asyncio
import logging
import shlex
from typing import Any, Optional

from server.application.server_batch_common import (
    agent_install_skip_result,
    batch_server_display_name,
    detect_and_save_target_path,
    ensure_agent_api_key,
    install_error_msg,
    mark_agent_uninstalled,
    prefetch_batch_context,
    result_item_dict,
    sudo_wrap,
)
from server.utils.agent_version import get_central_agent_version
from server.config import settings
from server.domain.models import AuditLog, Server
from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl
from server.infrastructure.database.server_repo import ServerRepositoryImpl
from server.infrastructure.database.session import AsyncSessionLocal
from server.infrastructure.redis import server_batch_store as sbs
from server.infrastructure.redis.client import get_redis
from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command

logger = logging.getLogger("nexus.server_batch_service")

_BG_TASKS: set[asyncio.Task] = set()
CONCURRENCY = 5

OP_LABELS = {
    "category": "批量修改分类",
    "health-check": "健康检查",
    "detect-path": "自动检测目标路径",
    "onboard": "新服务器初始化",
    "install-agent": "批量安装 Agent",
    "upgrade-agent": "批量升级 Agent",
    "uninstall-agent": "批量卸载 Agent",
}

AUDIT_ACTIONS = {
    "category": "batch_update_category",
    "health-check": "batch_health_check",
    "detect-path": "batch_detect_path",
    "onboard": "batch_server_onboard",
    "install-agent": "batch_install_agent",
    "upgrade-agent": "batch_upgrade_agent",
    "uninstall-agent": "batch_uninstall_agent",
}

HEARTBEAT_PREFIX = "heartbeat:"


async def _notify_progress(live: dict[str, Any]) -> None:
    try:
        from server.api.websocket import broadcast_script_progress

        stats = sbs.progress_stats(live)
        await broadcast_script_progress({
            "execution_id": int(live["id"]),
            "task_kind": "server_batch",
            "op": live.get("op") or "",
            "label": live.get("label") or "",
            "status": live.get("status") or "running",
            "long_task": False,
            "operator": live.get("operator") or "",
            **stats,
        })
    except Exception as e:
        logger.warning("server batch progress notify failed: %s", e)


async def _notify_complete(live: dict[str, Any]) -> None:
    try:
        from server.api.websocket import broadcast_script_complete

        stats = sbs.progress_stats(live)
        await broadcast_script_complete({
            "execution_id": int(live["id"]),
            "task_kind": "server_batch",
            "op": live.get("op") or "",
            "label": live.get("label") or "",
            "status": live.get("status") or "completed",
            "long_task": False,
            "operator": live.get("operator") or "",
            **stats,
        })
    except Exception as e:
        logger.warning("server batch complete notify failed: %s", e)


async def _reconcile_live_job(
    live: dict[str, Any],
    *,
    error: str,
    auto_stuck: bool,
) -> bool:
    """Fill pending servers as failed, finalize, persist, audit, notify. Returns True if reconciled."""
    if live.get("status") != "running":
        return False
    pending = sbs.pending_server_ids(live)
    if not pending and sbs.progress_stats(live)["remaining"] == 0:
        sbs.finalize_status(live)
        await sbs.save_live(live)
    else:
        added = sbs.apply_pending_failures(
            live,
            server_ids=pending,
            error=error,
            auto_stuck=auto_stuck,
        )
        if added == 0 and not pending:
            return False
        sbs.finalize_status(live)
        await sbs.save_live(live)
    await _persist_job_finalize(live)
    await _write_audit(live)
    await _notify_complete(live)
    job_id = int(live["id"])
    logger.info(
        "Server batch job reconciled job_id=%s op=%s status=%s auto_stuck=%s",
        job_id,
        live.get("op"),
        live.get("status"),
        auto_stuck,
    )
    return True


async def recover_orphaned_batch_jobs() -> int:
    """On API startup: finalize batch jobs left running after process restart."""
    from server.infrastructure.database.server_batch_job_repo import (
        ServerBatchJobRepositoryImpl,
    )

    reconciled_ids: set[int] = set()
    count = 0

    for live in await sbs.list_running_live_jobs():
        job_id = int(live["id"])
        if job_id in reconciled_ids:
            continue
        if await _reconcile_live_job(
            live,
            error=sbs.INTERRUPTED_MESSAGE,
            auto_stuck=False,
        ):
            reconciled_ids.add(job_id)
            count += 1

    async with AsyncSessionLocal() as session:
        repo = ServerBatchJobRepositoryImpl(session)
        for row in await repo.list_running():
            job_id = int(row.id)
            if job_id in reconciled_ids:
                continue
            live = await sbs.get_live(job_id)
            if live is None:
                live = sbs.live_from_mysql_row(row)
            elif live.get("status") != "running":
                continue
            if await _reconcile_live_job(
                live,
                error=sbs.INTERRUPTED_MESSAGE,
                auto_stuck=False,
            ):
                reconciled_ids.add(job_id)
                count += 1

    if count:
        logger.info("Recovered %s orphaned server batch job(s) after startup", count)
    return count


async def detect_stuck_batch_jobs() -> int:
    """Mark running batch jobs with no progress for STUCK_NO_PROGRESS_SECONDS."""
    count = 0
    for live in await sbs.list_running_live_jobs():
        if not sbs.is_stale_running(live):
            continue
        if await _reconcile_live_job(
            live,
            error=sbs.STUCK_MESSAGE,
            auto_stuck=True,
        ):
            count += 1
    if count:
        logger.info("Auto-reconciled %s stuck server batch job(s)", count)
    return count


async def start_batch_job(
    *,
    op: str,
    server_ids: list[int],
    operator: str,
    params: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    if op not in OP_LABELS:
        raise ValueError(f"未知的批量操作: {op}")
    if not server_ids:
        raise ValueError("server_ids 不能为空")

    job_id = await sbs.next_job_id()
    label = OP_LABELS[op]
    if op == "category" and params and params.get("category"):
        label = f"{label} → {params['category']}"

    live = sbs.init_live(
        job_id=job_id,
        op=op,
        label=label,
        server_ids=server_ids,
        operator=operator,
        params=params,
    )
    await sbs.save_live(live)
    await _persist_job_start(live)
    await _notify_progress(live)

    task = asyncio.create_task(_run_background(job_id))
    _BG_TASKS.add(task)
    task.add_done_callback(_BG_TASKS.discard)

    stats = sbs.progress_stats(live)
    return {
        "job_id": job_id,
        "status": "running",
        "label": label,
        "op": op,
        **stats,
    }


async def get_batch_job(job_id: int) -> Optional[dict[str, Any]]:
    live = await sbs.get_live(job_id)
    if live:
        return sbs.live_to_api(live)
    return await _get_batch_job_from_mysql(job_id)


async def list_batch_jobs(
    *,
    limit: int = 50,
    offset: int = 0,
    op: Optional[str] = None,
) -> dict[str, Any]:
    from server.infrastructure.database.server_batch_job_repo import (
        ServerBatchJobRepositoryImpl,
        job_row_to_api,
    )

    async with AsyncSessionLocal() as session:
        repo = ServerBatchJobRepositoryImpl(session)
        rows, total = await repo.list_recent(limit=limit, offset=offset, op=op)
        items = [job_row_to_api(row) for row in rows]
    return {"items": items, "total": total, "limit": limit, "offset": offset}


async def _persist_job_start(live: dict[str, Any]) -> None:
    from server.infrastructure.database.server_batch_job_repo import ServerBatchJobRepositoryImpl

    try:
        async with AsyncSessionLocal() as session:
            repo = ServerBatchJobRepositoryImpl(session)
            await repo.upsert_running(
                job_id=int(live["id"]),
                op=str(live.get("op") or ""),
                label=str(live.get("label") or ""),
                server_ids=list(live.get("server_ids") or []),
                operator=str(live.get("operator") or ""),
                params=live.get("params") if isinstance(live.get("params"), dict) else {},
            )
    except Exception as e:
        logger.warning("Server batch MySQL create failed job_id=%s: %s", live.get("id"), e)


async def _persist_job_finalize(live: dict[str, Any]) -> None:
    from server.infrastructure.database.server_batch_job_repo import ServerBatchJobRepositoryImpl

    try:
        api = sbs.live_to_api(live)
        async with AsyncSessionLocal() as session:
            repo = ServerBatchJobRepositoryImpl(session)
            await repo.finalize(
                int(live["id"]),
                status=str(live.get("status") or "completed"),
                results=list(api.get("results") or []),
            )
    except Exception as e:
        logger.warning("Server batch MySQL finalize failed job_id=%s: %s", live.get("id"), e)


async def _get_batch_job_from_mysql(job_id: int) -> Optional[dict[str, Any]]:
    from server.infrastructure.database.server_batch_job_repo import (
        ServerBatchJobRepositoryImpl,
        job_row_to_api,
    )

    try:
        async with AsyncSessionLocal() as session:
            repo = ServerBatchJobRepositoryImpl(session)
            row = await repo.get_by_id(job_id)
            if not row:
                return None
            return job_row_to_api(row)
    except Exception as e:
        logger.warning("Server batch MySQL read failed job_id=%s: %s", job_id, e)
        return None


async def _run_background(job_id: int) -> None:
    live = await sbs.get_live(job_id)
    if not live:
        return
    op = live.get("op") or ""
    try:
        if op == "category":
            await _run_category(live)
        elif op == "health-check":
            await _run_health_check(live)
        elif op == "detect-path":
            await _run_detect_path(live)
        elif op == "onboard":
            await _run_onboard(live)
        elif op == "install-agent":
            await _run_install_agent(live)
        elif op == "upgrade-agent":
            await _run_upgrade_agent(live)
        elif op == "uninstall-agent":
            await _run_uninstall_agent(live)
        else:
            live["status"] = "failed"
            await sbs.save_live(live)
            return
        live = await sbs.get_live(job_id)
        if live:
            stats = sbs.progress_stats(live)
            if live.get("status") == "running" or stats["remaining"] == 0:
                sbs.finalize_status(live)
                await sbs.save_live(live)
    except Exception:
        logger.exception("Server batch job failed job_id=%s op=%s", job_id, op)
        live = await sbs.get_live(job_id)
        if live:
            live["status"] = "failed"
            await sbs.save_live(live)
    live = await sbs.get_live(job_id)
    if live:
        if live.get("status") == "running":
            sbs.finalize_status(live)
            await sbs.save_live(live)
        await _persist_job_finalize(live)
        await _write_audit(live)
        await _notify_complete(live)


async def _write_audit(live: dict[str, Any]) -> None:
    import json

    op = live.get("op") or ""
    action = AUDIT_ACTIONS.get(op, "batch_server_op")
    stats = sbs.progress_stats(live)
    api = sbs.live_to_api(live)
    failures = [
        {"server_name": (e.get("server_name") or "").strip() or "未知服务器", "error": (e.get("error") or "")[:300]}
        for e in (api.get("results") or [])
        if not e.get("success")
    ]
    summary_map = {
        "category": f"批量修改分类 更新{stats['success']}台",
        "health-check": f"批量健康检查: {stats['success']}/{stats['total']} 在线",
        "detect-path": f"批量检测路径: {stats['success']}/{stats['total']} 成功",
        "onboard": f"新服务器初始化: {stats['success']}/{stats['total']} 成功",
        "install-agent": f"批量安装Agent: {stats['success']}/{stats['total']} 成功",
        "upgrade-agent": f"批量升级Agent: {stats['success']}/{stats['total']} 成功",
        "uninstall-agent": f"批量卸载Agent: {stats['success']}/{stats['total']} 成功",
    }
    summary = summary_map.get(op, f"{live.get('label')} {stats['success']}/{stats['total']}")
    detail = json.dumps({
        "job_id": int(live["id"]),
        "summary": summary,
        "success": stats["success"],
        "failed": stats["failed"],
        "total": stats["total"],
        "failures": failures[:30],
    }, ensure_ascii=False)
    try:
        async with AsyncSessionLocal() as session:
            audit_repo = AuditLogRepositoryImpl(session)
            await audit_repo.create(AuditLog(
                admin_username=live.get("operator") or "admin",
                action=action,
                target_type="server_batch_job",
                target_id=int(live["id"]),
                detail=detail,
                ip_address="",
            ))
            await session.commit()
    except Exception as e:
        logger.warning("Server batch audit failed job_id=%s: %s", live.get("id"), e)


async def _finish_one(live: dict[str, Any], item: dict[str, Any]) -> None:
    sbs.record_result(live, item)
    await sbs.save_live(live)
    await _notify_progress(live)


async def _run_with_semaphore(live: dict[str, Any], handler) -> None:
    from server.application.services.server_service import ServerService
    from server.infrastructure.database.server_repo import ServerRepositoryImpl
    from server.infrastructure.database.sync_log_repo import SyncLogRepositoryImpl

    server_ids = live.get("server_ids") or []
    sem = asyncio.Semaphore(CONCURRENCY)
    async with AsyncSessionLocal() as session:
        service = ServerService(
            server_repo=ServerRepositoryImpl(session),
            sync_log_repo=SyncLogRepositoryImpl(session),
        )
        server_map, labels = await prefetch_batch_context(service, server_ids)

        async def _one(sid: int) -> None:
            async with sem:
                item = await handler(sid, server_map, labels, live)
                await _finish_one(live, item)

        await asyncio.gather(*[_one(sid) for sid in server_ids])


async def _run_category(live: dict[str, Any]) -> None:
    params = live.get("params") or {}
    category = (params.get("category") or "").strip() or None
    server_ids = live.get("server_ids") or []

    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        servers = await repo.get_by_ids(server_ids)
        found_ids = {s.id for s in servers}
        labels = {s.id: batch_server_display_name(s) for s in servers}
        for sid in server_ids:
            labels.setdefault(sid, "未知服务器")

        for sid in server_ids:
            if sid not in found_ids:
                await _finish_one(live, result_item_dict(
                    server_id=sid,
                    server_name=labels[sid],
                    success=False,
                    error="服务器不存在",
                ))

        if found_ids:
            await repo.update_category_batch(list(found_ids), category)
            await session.commit()
            cat_label = category or "(清除)"
            for sid in found_ids:
                await _finish_one(live, result_item_dict(
                    server_id=sid,
                    server_name=labels[sid],
                    success=True,
                    stdout=f"分类 → {cat_label}",
                ))


async def _run_health_check(live: dict[str, Any]) -> None:
    from server.infrastructure.ssh.remote_probe import ssh_health_probe

    async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict:
        server = server_map.get(sid)
        label = labels[sid]
        if not server:
            return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在")
        try:
            probe = await ssh_health_probe(server, timeout=15)
            online = probe.get("status") == "online"
            detail = probe.get("error") or probe.get("channel") or probe.get("status") or ""
            return result_item_dict(
                server_id=sid,
                server_name=label,
                success=online,
                stdout=str(detail)[:500] if online else "",
                error="" if online else str(probe.get("error") or "离线")[:300],
            )
        except Exception as e:
            return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300])

    await _run_with_semaphore(live, handler)


async def _run_detect_path(live: dict[str, Any]) -> None:
    async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict:
        server = server_map.get(sid)
        label = labels[sid]
        if not server:
            return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在")
        try:
            outcome = await detect_and_save_target_path(server)
            if outcome.get("success"):
                return result_item_dict(
                    server_id=sid,
                    server_name=label,
                    success=True,
                    stdout=str(outcome.get("stdout") or ""),
                )
            return result_item_dict(
                server_id=sid,
                server_name=label,
                success=False,
                error=str(outcome.get("error") or "探测失败"),
            )
        except Exception as e:
            return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300])

    await _run_with_semaphore(live, handler)


async def _run_onboard(live: dict[str, Any]) -> None:
    from server.application.services.files_sudoers_service import install_nexus_files_sudoers
    from server.utils.posix_paths import is_unset_target_path

    async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict:
        server = server_map.get(sid)
        label = labels[sid]
        if not server:
            return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在")
        steps: list[str] = []
        try:
            ssh_user = (server.username or "root").strip() or "root"
            if ssh_user == "root":
                steps.append("sudo: skip(root)")
            else:
                sudo_outcome = await install_nexus_files_sudoers(server)
                action = sudo_outcome.get("action") or "unknown"
                if not sudo_outcome.get("ok"):
                    detail = sudo_outcome.get("detail") or action
                    return result_item_dict(
                        server_id=sid,
                        server_name=label,
                        success=False,
                        error=f"sudo 配置失败: {detail}",
                        stdout="; ".join(steps),
                    )
                steps.append(f"sudo: {action}")

            if is_unset_target_path(server.target_path):
                detect_outcome = await detect_and_save_target_path(server)
                if detect_outcome.get("success"):
                    steps.append(str(detect_outcome.get("stdout") or "path: ok"))
                else:
                    steps.append(f"path: {detect_outcome.get('error') or 'fail'}（推送仍可用 /www/wwwroot）")
            else:
                steps.append("path: already_set")

            return result_item_dict(
                server_id=sid,
                server_name=label,
                success=True,
                stdout="; ".join(steps),
            )
        except Exception as e:
            return result_item_dict(
                server_id=sid,
                server_name=label,
                success=False,
                error=str(e)[:300],
                stdout="; ".join(steps),
            )

    await _run_with_semaphore(live, handler)


async def _run_install_agent(live: dict[str, Any]) -> None:
    base_url = (settings.API_BASE_URL or "").strip().rstrip("/")
    if not base_url:
        for sid in live.get("server_ids") or []:
            await _finish_one(live, result_item_dict(
                server_id=sid, server_name=str(sid), success=False,
                error="NEXUS_API_BASE_URL 未配置",
            ))
        return

    central_version = get_central_agent_version()

    async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict:
        server = server_map.get(sid)
        label = labels[sid]
        if not server:
            return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在")
        try:
            skip = await agent_install_skip_result(
                server=server,
                server_id=sid,
                server_name=label,
                central_version=central_version,
            )
            if skip is not None:
                return skip
            api_key = await ensure_agent_api_key(server)
            ssh_user = (server.username or "root").strip()
            install_cmd = (
                f"TMP=$(mktemp) && curl -fsSL {shlex.quote(base_url)}/agent/install.sh -o \"$TMP\" && bash \"$TMP\" -- "
                f"--url {shlex.quote(base_url)} --key {shlex.quote(api_key)} "
                f"--id {sid} --port {shlex.quote(str(server.agent_port or 8601))}"
                f"; RC=$?; rm -f \"$TMP\"; exit $RC"
            )
            install_cmd = sudo_wrap(install_cmd, ssh_user)
            r = await exec_ssh_command(server, install_cmd, timeout=180)
            ok = r["exit_code"] == 0
            if ok and "Status: FAILED" in r.get("stdout", ""):
                ok = False
            return result_item_dict(
                server_id=sid, server_name=label, success=ok,
                stdout=r["stdout"][:2000] if ok else r["stdout"][:1000],
                error="" if ok else install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]),
            )
        except Exception as e:
            return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300])

    await _run_with_semaphore(live, handler)


async def _run_upgrade_agent(live: dict[str, Any]) -> None:
    base_url = (settings.API_BASE_URL or "").strip().rstrip("/")
    if not base_url:
        for sid in live.get("server_ids") or []:
            await _finish_one(live, result_item_dict(
                server_id=sid, server_name=str(sid), success=False,
                error="NEXUS_API_BASE_URL 未配置",
            ))
        return

    async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict:
        server = server_map.get(sid)
        label = labels[sid]
        if not server:
            return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在")
        try:
            from server.utils.agent_deploy import (
                AGENT_INSTALL_DIR,
                build_agent_upgrade_cmd,
                build_agent_upgrade_rollback_cmd,
            )

            install_dir = AGENT_INSTALL_DIR
            ssh_user = (server.username or "root").strip()
            agent_port = server.agent_port or 8601
            upgrade_cmd = build_agent_upgrade_cmd(
                base_url=base_url,
                ssh_user=ssh_user,
                install_dir=install_dir,
                agent_port=agent_port,
            )
            r = await exec_ssh_command(server, upgrade_cmd, timeout=120)
            ok = r["exit_code"] == 0 and "upgrade_ok" in r["stdout"]
            if not ok:
                try:
                    rollback = build_agent_upgrade_rollback_cmd(
                        ssh_user=ssh_user,
                        install_dir=install_dir,
                    )
                    await exec_ssh_command(server, rollback, timeout=30)
                except Exception as rb_err:
                    logger.warning("Rollback failed for server %s: %s", sid, rb_err)
            return result_item_dict(
                server_id=sid, server_name=label, success=ok,
                stdout=r["stdout"][:500] if ok else "",
                error=install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]) if not ok else "",
            )
        except Exception as e:
            return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:200])

    await _run_with_semaphore(live, handler)


async def _run_uninstall_agent(live: dict[str, Any]) -> None:
    base_url = (settings.API_BASE_URL or "").strip().rstrip("/")
    if not base_url:
        for sid in live.get("server_ids") or []:
            await _finish_one(live, result_item_dict(
                server_id=sid, server_name=str(sid), success=False,
                error="NEXUS_API_BASE_URL 未配置",
            ))
        return

    async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict:
        server = server_map.get(sid)
        label = labels[sid]
        if not server:
            return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在")
        try:
            ssh_user = (server.username or "root").strip()
            uninstall_cmd = (
                f"TMP=$(mktemp) && curl -fsSL {shlex.quote(base_url)}/agent/uninstall.sh -o \"$TMP\" "
                f"&& bash \"$TMP\" ; RC=$?; rm -f \"$TMP\"; exit $RC"
            )
            uninstall_cmd = sudo_wrap(uninstall_cmd, ssh_user)
            r = await exec_ssh_command(server, uninstall_cmd, timeout=60)
            ok = r["exit_code"] == 0
            if ok:
                try:
                    redis = get_redis()
                    await redis.delete(f"{HEARTBEAT_PREFIX}{sid}")
                except Exception as redis_err:
                    logger.warning("Failed to clear Redis heartbeat for server %s: %s", sid, redis_err)
                await mark_agent_uninstalled(sid)
            return result_item_dict(
                server_id=sid, server_name=label, success=ok,
                stdout=r["stdout"][:500] if ok else "",
                error="" if ok else install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]),
            )
        except Exception as e:
            return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300])

    await _run_with_semaphore(live, handler)