live-hub-py/web/backend/routers/huya/workbench.py

"""虎牙工作台：配置、商品快照、任务批次与日志。"""

from fastapi import APIRouter

from ._common import (
    HUYA_CONFIG_FIELDS,
    SUPPORTED_TASK_TYPES,
    UTC,
    Depends,
    HTTPException,
    HuyaBatchRunner,
    HuyaConfigOut,
    HuyaConfigUpdate,
    HuyaGoodsOut,
    HuyaGoodsSnapshot,
    HuyaRechargeGoodsOut,
    HuyaRechargeGoodsSnapshot,
    HuyaTask,
    HuyaTaskBatchRequest,
    HuyaTaskOut,
    HuyaWorkbench,
    HuyaWorkbenchAccount,
    HuyaWorkbenchAccountsUpdate,
    Query,
    Session,
    SessionLocal,
    User,
    WebSocket,
    WebSocketDisconnect,
    _config_out,
    _has_huya_perm,
    _huya_task_summary,
    _require_huya_batch_owner,
    _require_huya_task_account_access,
    _task_out,
    _visible_huya_tasks_query,
    _workbench_account_ids,
    apply_huya_config_defaults,
    asyncio,
    authenticate_websocket,
    cleanup_orphan_huya_tasks,
    create_planned_tasks,
    datetime,
    ensure_huya_config,
    get_db,
    huya_batch_registry,
    record_audit,
    require_permission,
    threading,
)

router = APIRouter()


@router.get("/task-types")
def task_types(current: User = Depends(require_permission("huya:task"))):
    """返回当前规划的虎牙任务类型。"""
    return SUPPORTED_TASK_TYPES


@router.get("/config", response_model=HuyaConfigOut)
def get_config(
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:config")),
):
    """获取虎牙配置。"""
    config = ensure_huya_config(db)
    return _config_out(config)


@router.put("/config", response_model=HuyaConfigOut)
def update_config(
    req: HuyaConfigUpdate,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:config")),
):
    """更新虎牙配置。"""
    config = ensure_huya_config(db)
    changed_fields = []
    for field in HUYA_CONFIG_FIELDS:
        value = getattr(req, field)
        if value is not None:
            setattr(config, field, value.strip())
            changed_fields.append(field)
    apply_huya_config_defaults(config)
    config.updated_at = datetime.now(UTC)
    if changed_fields:
        record_audit(
            db,
            current,
            action="recharge:huya:config",
            target="huya_config",
            detail={"changed_fields": changed_fields},
        )
    db.commit()
    db.refresh(config)
    return _config_out(config)


@router.get("/workbench-accounts")
def list_huya_workbench_accounts(
    handbook_scope: str = Query(..., pattern="^elite$|^browser_exchange$"),
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """读取当前用户指定工作台（精英宝典/浏览器兑换）的账号集合。"""
    ids = _workbench_account_ids(db, current, handbook_scope)
    configured = (
        db.query(HuyaWorkbench.id)
        .filter(
            HuyaWorkbench.user_id == current.id,
            HuyaWorkbench.handbook_scope == handbook_scope,
        )
        .first()
        is not None
    )
    return {"account_ids": ids, "configured": configured}


@router.put("/workbench-accounts")
def update_huya_workbench_accounts(
    req: HuyaWorkbenchAccountsUpdate,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """覆盖当前用户指定工作台（精英宝典/浏览器兑换）的账号集合，支持跨浏览器同步。"""
    ids = sorted(set(req.account_ids))
    if any(account_id < 1 for account_id in ids):
        raise HTTPException(status_code=400, detail="无效的账号 ID")
    if ids:
        _require_huya_task_account_access(db, current, ids)
    workbench = (
        db.query(HuyaWorkbench)
        .filter(
            HuyaWorkbench.user_id == current.id,
            HuyaWorkbench.handbook_scope == req.handbook_scope,
        )
        .first()
    )
    if workbench is None:
        db.add(HuyaWorkbench(user_id=current.id, handbook_scope=req.handbook_scope))
    else:
        workbench.updated_at = datetime.now(UTC)
    db.query(HuyaWorkbenchAccount).filter(
        HuyaWorkbenchAccount.user_id == current.id,
        HuyaWorkbenchAccount.handbook_scope == req.handbook_scope,
    ).delete(synchronize_session=False)
    db.add_all(
        [
            HuyaWorkbenchAccount(
                user_id=current.id,
                handbook_scope=req.handbook_scope,
                account_id=account_id,
            )
            for account_id in ids
        ]
    )
    db.commit()
    return {"account_ids": ids, "success": True}


@router.get("/goods", response_model=list[HuyaGoodsOut])
def list_goods(
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """查看已缓存的虎牙商品快照。"""
    rows = db.query(HuyaGoodsSnapshot).order_by(HuyaGoodsSnapshot.id.desc()).all()
    latest = {}
    for row in rows:
        latest.setdefault((row.sid, row.product_id), row)
    return list(latest.values())


@router.get("/recharge-goods", response_model=list[HuyaRechargeGoodsOut])
def list_recharge_goods(
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """查看已缓存的虎牙充值商品快照。"""
    rows = (
        db.query(HuyaRechargeGoodsSnapshot)
        .order_by(HuyaRechargeGoodsSnapshot.id.desc())
        .all()
    )
    latest = {}
    for row in rows:
        latest.setdefault((row.sid, row.spu_id), row)
    return list(latest.values())


@router.post("/tasks/batch")
async def create_task_batch(
    req: HuyaTaskBatchRequest,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """创建虎牙任务记录并启动后台执行器。"""
    if not req.account_ids:
        raise HTTPException(status_code=400, detail="请选择虎牙账号")
    _require_huya_task_account_access(db, current, req.account_ids)

    # 先清掉没有执行器的历史 running/pending，避免前端被假活跃批次锁死。
    # 不清理 planned：避免与刚创建的新任务产生竞态。
    cleanup_orphan_huya_tasks(
        db,
        active_batch_ids=huya_batch_registry.active_ids(),
        statuses=("pending", "running"),
        message="任务已中断（无执行器接管）",
    )

    try:
        batch_id, count = create_planned_tasks(
            db,
            req.account_ids,
            req.task_type,
            current.id,
            req.payload,
            req.handbook_scope,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    if count == 0:
        raise HTTPException(status_code=400, detail="没有有效的虎牙账号")

    if req.task_type == "create_recharge_order":
        record_audit(
            db,
            current,
            action="recharge:huya:create",
            target=f"huya_batch:{batch_id}",
            detail={
                "batch_id": batch_id,
                "task_type": req.task_type,
                "count": count,
                "account_count": len(req.account_ids),
            },
        )
        db.commit()

    log_queue = asyncio.Queue()
    loop = asyncio.get_running_loop()
    thread_db = SessionLocal()
    runner = HuyaBatchRunner(
        db=thread_db,
        batch_id=batch_id,
        task_type=req.task_type,
        payload=req.payload,
        log_queue=log_queue,
        loop=loop,
        concurrency=req.concurrency,
    )
    huya_batch_registry.register(batch_id, log_queue, loop, runner)

    thread = threading.Thread(target=runner.run, daemon=True)
    thread.start()

    return {"batch_id": batch_id, "count": count, "success": True}


@router.get("/tasks", response_model=list[HuyaTaskOut])
def list_tasks(
    batch_id: str | None = None,
    handbook_scope: str | None = Query(None, pattern="^(legacy|elite)$"),
    include_images: bool = Query(
        False, description="是否返回 base64 小程序码（默认否，轮询请保持 false）"
    ),
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """查看虎牙任务记录。"""
    query = _visible_huya_tasks_query(db, current)
    if batch_id:
        query = query.filter(HuyaTask.batch_id == batch_id)
    if handbook_scope:
        query = query.filter(HuyaTask.handbook_scope == handbook_scope)
    tasks = query.order_by(HuyaTask.id.desc()).limit(300).all()
    return [_task_out(task, include_images=include_images) for task in tasks]


@router.post("/tasks/cleanup-orphans")
def cleanup_orphan_tasks(
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """手动清理没有内存执行器接管的虎牙任务。"""
    cleaned = cleanup_orphan_huya_tasks(
        db,
        active_batch_ids=huya_batch_registry.active_ids(),
        statuses=("pending", "running"),
        message="任务已中断（无执行器接管）",
    )
    return {
        "message": f"已清理 {cleaned} 个残留任务",
        "cleaned": cleaned,
        "success": True,
    }


@router.get("/tasks/summary")
def tasks_summary(
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """虎牙任务统计。"""
    return _huya_task_summary(_visible_huya_tasks_query(db, current))


@router.get("/tasks/{task_id}", response_model=HuyaTaskOut)
def get_task(
    task_id: int,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """获取单条虎牙任务详情（含二维码图片）。"""
    task = _visible_huya_tasks_query(db, current).filter(HuyaTask.id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")
    return _task_out(task, include_images=True)


@router.post("/stop/{batch_id}")
def stop_batch(
    batch_id: str,
    db: Session = Depends(get_db),
    current: User = Depends(require_permission("huya:task")),
):
    """停止正在运行的虎牙批次。"""
    _require_huya_batch_owner(db, current, batch_id)
    is_recharge_batch = (
        db.query(HuyaTask.id)
        .filter(
            HuyaTask.batch_id == batch_id, HuyaTask.task_type == "create_recharge_order"
        )
        .first()
        is not None
    )
    batch = huya_batch_registry.get(batch_id)
    if batch:
        if batch.get("finished"):
            huya_batch_registry.pop(batch_id)
            cleaned = cleanup_orphan_huya_tasks(
                db,
                batch_id=batch_id,
                message="批次已结束",
            )
            if cleaned:
                return {
                    "message": f"批次已结束，已清理 {cleaned} 个残留任务",
                    "success": True,
                }
            raise HTTPException(status_code=404, detail="批次已结束")
        batch["runner"].stop()
        if is_recharge_batch:
            record_audit(
                db,
                current,
                action="recharge:huya:stop",
                target=f"huya_batch:{batch_id}",
                detail={"batch_id": batch_id, "mode": "running"},
            )
            db.commit()
        return {"message": "已发送停止信号", "success": True}

    cleaned = cleanup_orphan_huya_tasks(
        db,
        batch_id=batch_id,
        message="任务已停止（批次不存在）",
    )
    if cleaned:
        if is_recharge_batch:
            record_audit(
                db,
                current,
                action="recharge:huya:stop",
                target=f"huya_batch:{batch_id}",
                detail={
                    "batch_id": batch_id,
                    "mode": "orphan_cleanup",
                    "cleaned": cleaned,
                },
            )
            db.commit()
        return {"message": f"已清理 {cleaned} 个残留任务", "success": True}
    raise HTTPException(status_code=404, detail="批次不存在或已结束")


@router.websocket("/ws/{batch_id}")
async def ws_huya_logs(websocket: WebSocket, batch_id: str):
    """虎牙实时日志推送通道。"""
    user = authenticate_websocket(websocket)
    if not user:
        await websocket.close(code=1008, reason="未授权")
        return
    if not _has_huya_perm(user, "huya:task"):
        await websocket.close(code=1008, reason="无权限")
        return
    db = SessionLocal()
    try:
        _require_huya_batch_owner(db, user, batch_id)
    except HTTPException:
        await websocket.close(code=1008, reason="无权访问该任务批次")
        return
    finally:
        db.close()
    await websocket.accept()

    batch = huya_batch_registry.get(batch_id)
    if not batch:
        await websocket.send_json({"level": "error", "message": "批次不存在或已结束"})
        await websocket.close()
        return

    log_queue: asyncio.Queue = batch["log_queue"]

    try:
        while True:
            try:
                msg = await asyncio.wait_for(log_queue.get(), timeout=30)
                await websocket.send_json(msg)
                if msg.get("level") == "result":
                    await asyncio.sleep(0.1)
                    break
            except TimeoutError:
                await websocket.send_json({"level": "heartbeat", "message": ""})
    except WebSocketDisconnect:
        pass
    finally:
        latest = huya_batch_registry.get(batch_id)
        if latest and latest.get("finished"):
            huya_batch_registry.pop(batch_id)