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 (
    batch_server_display_name,
    ensure_agent_api_key,
    install_error_msg,
    mark_agent_uninstalled,
    posix_dirname,
    prefetch_batch_context,
    result_item_dict,
    sudo_wrap,
    update_server_target_path,
)
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": "自动检测目标路径",
    "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",
    "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 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"Unknown batch op: {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 == "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 and live.get("status") == "running":
            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:
        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']} 成功",
        "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:
            ssh_user = (server.username or "root").strip()
            cmd = "find /www/wwwroot -iname 'workerman.bat' -type f -print -quit 2>/dev/null"
            cmd = sudo_wrap(cmd, ssh_user)
            r = await exec_ssh_command(server, cmd, timeout=30)
            if r["exit_code"] != 0 and not r.get("stdout", "").strip():
                return result_item_dict(
                    server_id=sid, server_name=label, success=False,
                    error=f"搜索失败 (exit {r['exit_code']})",
                )
            found = r["stdout"].strip()
            if not found:
                return result_item_dict(server_id=sid, server_name=label, success=False, error="未找到 workerman.bat")
            first_match = found.split("\n")[0].strip()
            target_dir = posix_dirname(first_match)
            await update_server_target_path(sid, target_dir)
            return result_item_dict(
                server_id=sid, server_name=label, success=True,
                stdout=f"target_path → {target_dir}",
            )
        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_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

    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:
            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:
            agent_url = f"{base_url}/agent/agent.py"
            install_dir = "/opt/nexus-agent"
            venv_python = f"{install_dir}/.venv/bin/python"
            ssh_user = (server.username or "root").strip()
            batch_systemctl_prefix = "sudo " if ssh_user != "root" else ""
            batch_kill_port = "fuser -k 8601/tcp 2>/dev/null || true; "
            upgrade_cmd = sudo_wrap(
                f"(command -v rsync >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq rsync) || (yum install -y -q rsync) || (dnf install -y -q rsync)) "
                f"&& {venv_python} -c 'import sys; v=sys.version_info; sys.exit(0 if v.major==3 and v.minor>=10 else 1)' "
                f"&& {venv_python} -m pip install -q fastapi==0.115.6 uvicorn==0.34.0 httpx==0.28.1 psutil python-multipart==0.0.19 "
                f"&& cp {install_dir}/agent.py {install_dir}/agent.py.bak "
                f"&& curl -fsSL {shlex.quote(agent_url)} -o {install_dir}/agent.py "
                f"&& {batch_kill_port}{batch_systemctl_prefix}systemctl restart nexus-agent "
                f"&& sleep 2 "
                f"&& {batch_systemctl_prefix}systemctl is-active nexus-agent "
                f"&& echo 'upgrade_ok'",
                ssh_user,
            )
            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 = sudo_wrap(
                        f"cp {install_dir}/agent.py.bak {install_dir}/agent.py && {batch_systemctl_prefix}systemctl restart nexus-agent",
                        ssh_user,
                    )
                    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)