live-hub-py/web/backend/services/huya_runner_core.py

"""虎牙任务执行器：公共基础（由 huya_runner.py 按功能域拆分）。"""

from __future__ import annotations

import asyncio
import threading
import time
from datetime import datetime, timezone
from typing import Optional

from loguru import logger
from sqlalchemy.orm import Session
from ..models import HuyaTask
from .huya_service import cookie_value

# 绑定跳转与充值下单共用（bind 与 recharge 两个域都引用了 sourceId）
HUYA_RECHARGE_SOURCE_ID = "yellowcarlist"
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .huya_runner import HuyaBatchRunner

class HuyaBatchRunnerCore:
    """虎牙任务执行器公共基础：批次状态、日志、任务落库。"""
    def __init__(
        self,
        db: Session,
        batch_id: str,
        task_type: str,
        payload: Optional[dict] = None,
        log_queue: Optional[asyncio.Queue] = None,
        loop: Optional[asyncio.AbstractEventLoop] = None,
        concurrency: int = 3,
    ):
        self.db = db
        self.batch_id = batch_id
        self.task_type = task_type
        self.payload = payload or {}
        self.log_queue = log_queue
        self.loop = loop
        self.concurrency = max(1, min(concurrency, 10))
        self._stop = threading.Event()
        self._counter_lock = threading.Lock()
        self._started = 0

    def stop(self):
        self._stop.set()

    def _push_log(self, level: str, message: str):
        if level == "result":
            try:
                huya_batch_registry.mark_finished(self.batch_id)
            except NameError:
                pass
        if level != "result" and message:
            log_func = getattr(logger, level, logger.info)
            log_func(f"[huya] {message}")
        if self.log_queue and self.loop:
            asyncio.run_coroutine_threadsafe(
                self.log_queue.put({"level": level, "message": message}),
                self.loop,
            )

    @staticmethod
    def _account_name(account_info: dict) -> str:
        return (
            account_info.get("nickname")
            or account_info.get("username")
            or account_info.get("uid")
            or f"#{account_info.get('account_id')}"
        )

    @staticmethod
    def _to_int(value) -> int:
        text = str(value or "").strip()
        return int(text) if text.isdigit() else 0

    def _resolve_uid(self, account_info: dict) -> int:
        cookie = account_info.get("cookie") or ""
        return (
            self._to_int(account_info.get("yyuid"))
            or self._to_int(account_info.get("uid"))
            or self._to_int(cookie_value(cookie, "yyuid"))
            or self._to_int(cookie_value(cookie, "udb_uid"))
        )

    @staticmethod
    def _format_local_time(timestamp: int) -> str:
        if not timestamp:
            return ""
        return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S")

    @staticmethod
    def _parse_scheduled_time(value) -> datetime | None:
        text = str(value or "").strip()
        if not text:
            return None
        try:
            normalized = text.replace("Z", "+00:00")
            dt = datetime.fromisoformat(normalized)
        except ValueError:
            return None
        if dt.tzinfo is None:
            return dt.astimezone()
        return dt

    def _wait_until(self, when: datetime, uid: int) -> bool:
        target = when.timestamp()
        local_text = self._format_local_time(int(target))
        self._push_log("info", f"[{uid}] 定时兑换等待到 {local_text}")
        while not self._stop.is_set():
            remaining = target - time.time()
            if remaining <= 0:
                return True
            time.sleep(min(0.2, max(0.02, remaining)))
        return False

    def _mark_task(
        self,
        worker_db: Session,
        task: HuyaTask,
        status: str,
        message: str,
        result: Optional[dict] = None,
    ):
        task.status = status
        task.message = message
        task.result = result
        task.finished_at = datetime.now(timezone.utc)
        worker_db.commit()

    def _update_task_progress(
        self,
        worker_db: Session,
        task: HuyaTask,
        status: str,
        message: str,
        result: Optional[dict] = None,
    ):
        task.status = status
        task.message = message
        if result is not None:
            task.result = result
        worker_db.commit()


class HuyaBatchRegistry:
    """管理运行中的虎牙批次。"""

    def __init__(self):
        self._batches: dict[str, dict] = {}
        self._lock = threading.Lock()

    def _cleanup_locked(self, ttl_seconds: int = 300):
        now = time.time()
        expired = [
            batch_id
            for batch_id, batch in self._batches.items()
            if batch.get("finished") and now - float(batch.get("finished_at") or now) > ttl_seconds
        ]
        for batch_id in expired:
            self._batches.pop(batch_id, None)

    def register(
        self,
        batch_id: str,
        log_queue: asyncio.Queue,
        loop: asyncio.AbstractEventLoop,
        runner: HuyaBatchRunner,
    ):
        with self._lock:
            self._cleanup_locked()
            self._batches[batch_id] = {
                "log_queue": log_queue,
                "loop": loop,
                "runner": runner,
                "finished": False,
                "finished_at": None,
            }

    def get(self, batch_id: str):
        with self._lock:
            self._cleanup_locked()
            return self._batches.get(batch_id)

    def active_ids(self) -> set[str]:
        with self._lock:
            self._cleanup_locked()
            return {
                batch_id
                for batch_id, batch in self._batches.items()
                if not batch.get("finished")
            }

    def mark_finished(self, batch_id: str):
        with self._lock:
            batch = self._batches.get(batch_id)
            if not batch:
                return
            batch["finished"] = True
            batch["finished_at"] = time.time()

    def pop(self, batch_id: str):
        with self._lock:
            return self._batches.pop(batch_id, None)


huya_batch_registry = HuyaBatchRegistry()