core/task_manager.py

from __future__ import annotations

import logging
import re
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime
from zoneinfo import ZoneInfo

from core.course_bot import CourseBot, TaskResult
from core.db import Database
from core.login_modes import LOGIN_MODE_LABELS, LOGIN_MODE_UNIFIED, normalize_login_mode
from core.security import SecretBox
from core.task_modes import TASK_RUN_MODE_LABELS, TASK_RUN_MODE_STABLE, normalize_task_run_mode


LOGGER = logging.getLogger("sacc.task_manager")


@dataclass(slots=True)
class RunningTask:
    task_id: int
    thread: threading.Thread
    stop_event: threading.Event


@dataclass(slots=True)
class LoginTwoFactorChallenge:
    task_id: int
    user_id: int
    student_id: str
    display_name: str
    phone_mask: str
    created_at: float
    expires_at: float
    event: threading.Event = field(default_factory=threading.Event)
    submitted_code: str = ""
    submitted_by: str = ""


class TaskManager:
    login_2fa_timeout_seconds = 300

    def __init__(self, *, config, store: Database, secret_box: SecretBox) -> None:
        self.config = config
        self.store = store
        self.secret_box = secret_box
        self._started = False
        self._startup_lock = threading.Lock()
        self._queue_lock = threading.Lock()
        self._running_lock = threading.Lock()
        self._shutdown_event = threading.Event()
        self._dispatcher_thread: threading.Thread | None = None
        self._running: dict[int, RunningTask] = {}
        self._last_dispatch_snapshot: tuple[int, int, int] | None = None
        self._schedule_timezone = ZoneInfo(getattr(self.config, "schedule_timezone", "Asia/Shanghai"))
        self._login_2fa_lock = threading.Lock()
        self._login_2fa_challenges: dict[int, LoginTwoFactorChallenge] = {}

    def start(self) -> None:
        with self._startup_lock:
            if self._started:
                LOGGER.info("Task manager start skipped because it is already running")
                return
            LOGGER.info(
                "Task manager starting | db_path=%s default_parallel_limit=%s schedule_timezone=%s",
                self.store.path,
                getattr(self.config, "default_parallel_limit", "-"),
                getattr(self.config, "schedule_timezone", "Asia/Shanghai"),
            )
            self.store.reset_inflight_tasks()
            normalized_modes = self.store.normalize_legacy_task_modes()
            if normalized_modes:
                LOGGER.info("Normalized %s legacy task mode record(s) to stable", normalized_modes)
            normalized_login_modes = self.store.normalize_legacy_login_modes()
            if normalized_login_modes:
                LOGGER.info("Normalized %s legacy login mode record(s) to unified", normalized_login_modes)
            self._dispatcher_thread = threading.Thread(target=self._dispatch_loop, name="task-dispatcher", daemon=True)
            self._dispatcher_thread.start()
            self._started = True
            LOGGER.info("Task manager started")

    def shutdown(self) -> None:
        LOGGER.info("Task manager shutdown requested")
        self._shutdown_event.set()
        with self._running_lock:
            for running_task in self._running.values():
                running_task.stop_event.set()
        with self._login_2fa_lock:
            self._login_2fa_challenges.clear()
        LOGGER.info("Task manager stop signal broadcast to %s running task(s)", self._running_count())

    def queue_task(
        self,
        user_id: int,
        requested_by: str,
        requested_by_role: str,
        requested_mode: str = TASK_RUN_MODE_STABLE,
        login_mode: str = LOGIN_MODE_UNIFIED,
        use_proxy: bool = False,
    ) -> tuple[dict, bool]:
        normalized_mode = normalize_task_run_mode(requested_mode)
        normalized_login_mode = normalize_login_mode(login_mode)
        with self._queue_lock:
            active_task = self.store.find_active_task_for_user(user_id)
            if active_task is None:
                task_id = self.store.create_task(
                    user_id,
                    requested_by,
                    requested_by_role,
                    requested_mode=normalized_mode,
                    login_mode=normalized_login_mode,
                    use_proxy=use_proxy,
                )
            else:
                task_id = int(active_task["id"])

        if active_task is not None:
            LOGGER.info(
                "Task queue skipped because an active task already exists | task_id=%s user_id=%s status=%s",
                active_task["id"],
                user_id,
                active_task["status"],
            )
            self._log(
                int(active_task["id"]),
                user_id,
                "SYSTEM",
                "INFO",
                f"收到新的启动请求，但已有活动任务处于 {active_task['status']} 状态，本次不重复创建。触发者: {requested_by_role}:{requested_by}。",
            )
            return self.store.get_task(active_task["id"]) or active_task, False

        course_count = len(self.store.list_courses_for_user(user_id))
        LOGGER.info(
            "Task queued | task_id=%s user_id=%s requested_by_role=%s requested_by=%s requested_mode=%s login_mode=%s use_proxy=%s",
            task_id,
            user_id,
            requested_by_role,
            requested_by,
            normalized_mode,
            normalized_login_mode,
            bool(use_proxy),
        )
        self._log(
            task_id,
            user_id,
            "SYSTEM",
            "INFO",
            f"任务已进入队列，触发者: {requested_by_role}:{requested_by}，模式: {self._mode_label(normalized_mode)}，登录方式: {self._login_mode_label(normalized_login_mode)}，代理: {'启用' if use_proxy else '关闭'}，待选课程 {course_count} 门。",
        )
        return self.store.get_task(task_id), True

    def stop_task(self, task_id: int) -> bool:
        requested = self.store.request_task_stop(task_id)
        if not requested:
            LOGGER.info("Task stop request ignored because task is not active | task_id=%s", task_id)
            return False
        LOGGER.info("Task stop requested | task_id=%s", task_id)
        self._log(task_id, None, "SYSTEM", "INFO", "收到停止请求，任务会在安全节点退出。")
        with self._running_lock:
            running_task = self._running.get(task_id)
            if running_task is not None:
                running_task.stop_event.set()
        self._remove_login_2fa_challenge(task_id)
        return True

    def request_login_2fa_code(self, task_id: int, user: dict, phone_mask: str, stop_event: threading.Event) -> str:
        now = time.time()
        challenge = LoginTwoFactorChallenge(
            task_id=task_id,
            user_id=int(user["id"]),
            student_id=str(user.get("student_id") or ""),
            display_name=str(user.get("display_name") or ""),
            phone_mask=str(phone_mask or "绑定手机"),
            created_at=now,
            expires_at=now + self.login_2fa_timeout_seconds,
        )
        with self._login_2fa_lock:
            self._prune_login_2fa_challenges_locked(now=now)
            self._login_2fa_challenges[task_id] = challenge

        self.store.update_task_status(task_id, "waiting_2fa", "等待短信验证码。")
        self._log(
            task_id,
            int(user["id"]),
            "SYSTEM",
            "INFO",
            f"已发送短信验证码，请在 {self.login_2fa_timeout_seconds // 60} 分钟内到面板输入 6 位验证码。手机号={challenge.phone_mask}。",
        )

        next_wait_log_at = now + 60
        while time.time() < challenge.expires_at:
            if stop_event.is_set() or self._shutdown_event.is_set():
                self._remove_login_2fa_challenge(task_id)
                self._log(task_id, int(user["id"]), "SYSTEM", "WARNING", "短信验证码等待被停止，已清理本次验证码挑战。")
                raise RuntimeError("短信验证码等待已停止。")
            if challenge.event.wait(timeout=0.2):
                code = challenge.submitted_code
                self._remove_login_2fa_challenge(task_id)
                if not code:
                    self._log(task_id, int(user["id"]), "SYSTEM", "WARNING", "短信验证码等待被停止，未收到有效提交。")
                    raise RuntimeError("短信验证码等待已停止。")
                current_task = self.store.get_task(task_id)
                if current_task and current_task.get("status") == "waiting_2fa":
                    self.store.update_task_status(task_id, "running", "")
                self._log(task_id, int(user["id"]), "SYSTEM", "INFO", "短信验证码已转交后台 Selenium，会继续完成统一认证。")
                return code
            current_time = time.time()
            if current_time >= next_wait_log_at:
                remaining_seconds = max(0, int(challenge.expires_at - current_time))
                self._log(
                    task_id,
                    int(user["id"]),
                    "SYSTEM",
                    "INFO",
                    f"仍在等待短信验证码提交，剩余约 {remaining_seconds} 秒。",
                )
                next_wait_log_at = current_time + 60

        self._remove_login_2fa_challenge(task_id)
        self._log(task_id, int(user["id"]), "SYSTEM", "ERROR", "等待短信验证码超时，已清理本次验证码挑战。")
        raise TimeoutError("等待短信验证码超时，请重新启动任务。")

    def submit_login_2fa_code(self, task_id: int, code: str, *, submitted_by: str) -> tuple[bool, str]:
        normalized_code = str(code or "").strip()
        if not re.fullmatch(r"\d{6}", normalized_code):
            return False, "短信验证码必须是 6 位数字。"
        now = time.time()
        with self._login_2fa_lock:
            self._prune_login_2fa_challenges_locked(now=now)
            challenge = self._login_2fa_challenges.get(task_id)
            if challenge is None:
                return False, "当前没有等待验证码的任务，或验证码已过期。"
            if challenge.event.is_set():
                return False, "验证码已经提交，请等待后台继续登录。"
            challenge.submitted_code = normalized_code
            challenge.submitted_by = submitted_by
            challenge.event.set()

        self._log(task_id, challenge.user_id, "SYSTEM", "INFO", f"已收到短信验证码，后台继续登录。提交者: {submitted_by}")
        return True, "验证码已提交，后台会继续完成登录。"

    def get_login_2fa_challenge(self, task_id: int) -> dict | None:
        with self._login_2fa_lock:
            self._prune_login_2fa_challenges_locked()
            challenge = self._login_2fa_challenges.get(task_id)
            return None if challenge is None else self._serialize_login_2fa_challenge(challenge)

    def get_login_2fa_challenge_for_user(self, user_id: int) -> dict | None:
        with self._login_2fa_lock:
            self._prune_login_2fa_challenges_locked()
            for challenge in self._login_2fa_challenges.values():
                if challenge.user_id == int(user_id):
                    return self._serialize_login_2fa_challenge(challenge)
        return None

    def list_login_2fa_challenges(self) -> list[dict]:
        with self._login_2fa_lock:
            self._prune_login_2fa_challenges_locked()
            challenges = [self._serialize_login_2fa_challenge(challenge) for challenge in self._login_2fa_challenges.values()]
        return sorted(challenges, key=lambda item: int(item["remaining_seconds"]))

    def _remove_login_2fa_challenge(self, task_id: int) -> None:
        with self._login_2fa_lock:
            self._login_2fa_challenges.pop(task_id, None)

    def _prune_login_2fa_challenges_locked(self, *, now: float | None = None) -> None:
        current_time = time.time() if now is None else now
        expired_task_ids = [
            task_id
            for task_id, challenge in self._login_2fa_challenges.items()
            if current_time >= challenge.expires_at
        ]
        for task_id in expired_task_ids:
            self._login_2fa_challenges.pop(task_id, None)

    @staticmethod
    def _serialize_login_2fa_challenge(challenge: LoginTwoFactorChallenge) -> dict:
        remaining_seconds = max(0, int(challenge.expires_at - time.time() + 0.999))
        return {
            "task_id": challenge.task_id,
            "user_id": challenge.user_id,
            "student_id": challenge.student_id,
            "display_name": challenge.display_name,
            "phone_mask": challenge.phone_mask,
            "remaining_seconds": remaining_seconds,
        }

    def _current_schedule_now(self) -> datetime:
        return datetime.now(self._schedule_timezone)

    def _within_schedule_date_window(self, schedule: dict, today_text: str) -> bool:
        start_date = str(schedule.get("start_date") or "").strip()
        end_date = str(schedule.get("end_date") or "").strip()
        if start_date and today_text < start_date:
            return False
        if end_date and today_text > end_date:
            return False
        return True

    def _apply_user_schedules(self) -> None:
        now = self._current_schedule_now()
        today_text = now.date().isoformat()
        current_time_text = now.strftime("%H:%M")

        for schedule in self.store.list_enabled_user_schedules():
            if not bool(schedule.get("user_is_active", 0)):
                continue
            if not self._within_schedule_date_window(schedule, today_text):
                continue

            start_time = str(schedule.get("daily_start_time") or "").strip()
            stop_time = str(schedule.get("daily_stop_time") or "").strip()
            if not start_time or not stop_time:
                continue

            user_id = int(schedule["user_id"])
            active_task = self.store.find_active_task_for_user(user_id)

            if current_time_text >= stop_time and str(schedule.get("last_auto_stop_on") or "") != today_text:
                if active_task and self.stop_task(int(active_task["id"])):
                    self._log(int(active_task["id"]), user_id, "SYSTEM", "INFO", "已按管理员定时设置发送停止请求。")
                else:
                    self._log(None, user_id, "SYSTEM", "INFO", "已到定时停止时间，当前没有运行中的任务。")
                self.store.mark_schedule_auto_stop(user_id, today_text)
                continue

            if not (start_time <= current_time_text < stop_time):
                continue
            if str(schedule.get("last_auto_start_on") or "") == today_text:
                continue

            user = self.store.get_user(user_id)
            if user is None:
                self.store.mark_schedule_auto_start(user_id, today_text)
                continue
            if not self.store.list_courses_for_user(user_id):
                self._log(None, user_id, "SYSTEM", "INFO", "已到定时启动时间，但当前没有课程目标，今天不自动启动。")
                self.store.mark_schedule_auto_start(user_id, today_text)
                continue

            task, created = self.queue_task(
                user_id,
                requested_by="scheduler",
                requested_by_role="system",
                requested_mode=TASK_RUN_MODE_STABLE,
                login_mode=LOGIN_MODE_UNIFIED,
                use_proxy=False,
            )
            if created:
                self._log(task["id"], user_id, "SYSTEM", "INFO", "已按管理员定时设置自动加入任务队列。")
            else:
                self._log(task["id"], user_id, "SYSTEM", "INFO", "已到定时启动时间，但当前已有任务在运行或排队。")
            self.store.mark_schedule_auto_start(user_id, today_text)

    def _dispatch_loop(self) -> None:
        LOGGER.info("Dispatcher loop started")
        while not self._shutdown_event.is_set():
            self._cleanup_running_registry()
            self._apply_user_schedules()
            parallel_limit = self.store.get_parallel_limit()
            running_count = self._running_count()
            available_slots = max(0, parallel_limit - running_count)
            pending_tasks: list[dict] = []
            if available_slots > 0:
                pending_tasks = self.store.list_pending_tasks(available_slots)

            snapshot = (parallel_limit, running_count, len(pending_tasks))
            if snapshot != self._last_dispatch_snapshot and (running_count > 0 or pending_tasks):
                LOGGER.info(
                    "Dispatcher snapshot | parallel_limit=%s running=%s available_slots=%s fetched_pending=%s",
                    parallel_limit,
                    running_count,
                    available_slots,
                    len(pending_tasks),
                )
                self._last_dispatch_snapshot = snapshot

            if pending_tasks:
                for task in pending_tasks:
                    if self._shutdown_event.is_set():
                        break
                    if not task["is_active"]:
                        LOGGER.warning(
                            "Queued task cancelled because user is inactive | task_id=%s user_id=%s",
                            task["id"],
                            task["user_id"],
                        )
                        self.store.finish_task(task["id"], "failed", "该用户已被禁用，任务未执行。")
                        self._log(task["id"], task["user_id"], "SYSTEM", "WARNING", "用户已禁用，队列中的任务被取消。")
                        continue
                    self._log(
                        task["id"],
                        task["user_id"],
                        "SYSTEM",
                        "INFO",
                        f"调度器已获取执行名额，准备启动任务线程。当前并发 {running_count}/{parallel_limit}。",
                    )
                    self._launch_task(task["id"])
            time.sleep(1)
        LOGGER.info("Dispatcher loop stopped")

    def _launch_task(self, task_id: int) -> None:
        with self._running_lock:
            if task_id in self._running:
                LOGGER.info("Task launch skipped because thread already exists | task_id=%s", task_id)
                return
            stop_event = threading.Event()
            thread = threading.Thread(target=self._run_task, args=(task_id, stop_event), name=f"task-{task_id}", daemon=True)
            self._running[task_id] = RunningTask(task_id=task_id, thread=thread, stop_event=stop_event)
            LOGGER.info("Launching task thread | task_id=%s thread_name=%s", task_id, thread.name)
            task = self.store.get_task(task_id)
            self._log(
                task_id,
                int(task["user_id"]) if task else None,
                "SYSTEM",
                "INFO",
                f"调度器已分配执行线程 {thread.name}，任务即将开始运行。",
            )
            thread.start()

    def _run_task(self, task_id: int, stop_event: threading.Event) -> None:
        task = self.store.get_task(task_id)
        if task is None:
            LOGGER.warning("Task record disappeared before execution | task_id=%s", task_id)
            self._remove_running(task_id)
            return

        user = self.store.get_user(task["user_id"])
        if user is None:
            LOGGER.error("Task cannot start because user does not exist | task_id=%s user_id=%s", task_id, task["user_id"])
            self.store.finish_task(task_id, "failed", "用户不存在。")
            self._remove_running(task_id)
            return

        LOGGER.info("Task runner starting | task_id=%s user_id=%s", task_id, user["id"])
        if not self.store.mark_task_running(task_id):
            latest_task = self.store.get_task(task_id) or task
            latest_status = str(latest_task.get("status") or "")
            if latest_status == "cancel_requested":
                self.store.finish_task(task_id, "stopped", "任务启动前已收到停止请求。")
                self._log(task_id, user["id"], "SYSTEM", "INFO", "任务启动前已收到停止请求，未启动 Selenium。")
            else:
                self._log(task_id, user["id"], "SYSTEM", "WARNING", f"任务启动前状态已变为 {latest_status or '未知'}，未覆盖为运行态。")
            self._remove_running(task_id)
            return
        course_count = len(self.store.list_courses_for_user(user["id"]))
        requested_mode = normalize_task_run_mode(task.get("requested_mode"))
        effective_mode = normalize_task_run_mode(task.get("effective_mode"))
        login_mode = normalize_login_mode(task.get("login_mode"))
        use_proxy = bool(int(task.get("use_proxy") or 0))
        refresh_interval = user.get("refresh_interval_seconds") or getattr(self.config, "poll_interval_seconds", "-")
        self._log(
            task_id,
            user["id"],
            "SYSTEM",
            "INFO",
            "任务开始执行。"
            f"学号={user.get('student_id') or '-'}，"
            f"待选课程 {course_count} 门，"
            f"请求模式={self._mode_label(requested_mode)}，当前模式={self._mode_label(effective_mode)}，"
            f"登录方式={self._login_mode_label(login_mode)}，"
            f"代理={'启用' if use_proxy else '关闭'}，"
            f"刷新间隔={refresh_interval} 秒。",
        )

        try:
            password = self.secret_box.decrypt(user["password_encrypted"])
            bot = CourseBot(
                config=self.config,
                store=self.store,
                task_id=task_id,
                user=user,
                password=password,
                logger=lambda level, message: self._log(task_id, user["id"], "RUNNER", level, message),
                login_2fa_code_provider=lambda phone_mask, current_stop_event: self.request_login_2fa_code(
                    task_id,
                    user,
                    phone_mask,
                    current_stop_event,
                ),
            )
            result = bot.run(stop_event)
        except Exception as exc:  # pragma: no cover - defensive fallback
            LOGGER.exception("Task initialization failed | task_id=%s user_id=%s", task_id, user["id"])
            result = TaskResult(status="failed", error=str(exc))
            self._log(task_id, user["id"], "SYSTEM", "ERROR", f"任务初始化失败: {exc}")

        current_task = self.store.get_task(task_id)
        final_status = result.status
        if current_task and current_task["status"] == "cancel_requested" and result.status != "completed":
            final_status = "stopped"
        elif stop_event.is_set() and result.status != "completed":
            final_status = "stopped"

        self.store.finish_task(task_id, final_status, result.error)
        final_task = self.store.get_task(task_id) or {}
        final_attempts = int(final_task.get("total_attempts") or 0)
        final_errors = int(final_task.get("total_errors") or 0)
        remaining_courses = len(self.store.list_courses_for_user(user["id"]))
        final_text = result.error or f"任务已结束，状态: {final_status}"
        final_text = f"{final_text} 总尝试 {final_attempts} 次，累计异常 {final_errors} 次，剩余待选课程 {remaining_courses} 门。"
        level = "ERROR" if final_status == "failed" else "INFO"
        self._log(task_id, user["id"], "SYSTEM", level, final_text)
        LOGGER.info(
            "Task runner finished | task_id=%s user_id=%s final_status=%s",
            task_id,
            user["id"],
            final_status,
        )
        self._remove_running(task_id)

    def _log(self, task_id: int | None, user_id: int | None, scope: str, level: str, message: str) -> None:
        upper_level = level.upper()
        self.store.add_log(task_id, user_id, scope, upper_level, message)
        if task_id is not None and upper_level in {"WARNING", "ERROR"}:
            self.store.increment_task_errors(task_id)
        self._emit_runtime_log(task_id=task_id, user_id=user_id, scope=scope, level=upper_level, message=message)

    @staticmethod
    def _emit_runtime_log(*, task_id: int | None, user_id: int | None, scope: str, level: str, message: str) -> None:
        upper_level = level.upper()
        rendered = f"task_id={task_id or '-'} user_id={user_id or '-'} scope={scope} level={upper_level} | {message}"
        if upper_level == "ERROR":
            LOGGER.error(rendered)
        elif upper_level == "WARNING":
            LOGGER.warning(rendered)
        else:
            LOGGER.info(rendered)

    def _running_count(self) -> int:
        with self._running_lock:
            return len(self._running)

    @staticmethod
    def _mode_label(mode: str) -> str:
        return TASK_RUN_MODE_LABELS.get(mode, mode)

    @staticmethod
    def _login_mode_label(mode: str) -> str:
        return LOGIN_MODE_LABELS.get(mode, mode)

    def _cleanup_running_registry(self) -> None:
        finished_ids: list[int] = []
        with self._running_lock:
            for task_id, running_task in self._running.items():
                if not running_task.thread.is_alive():
                    finished_ids.append(task_id)
            for task_id in finished_ids:
                self._running.pop(task_id, None)
        for task_id in finished_ids:
            LOGGER.info("Task thread cleaned from registry | task_id=%s", task_id)

    def _remove_running(self, task_id: int) -> None:
        with self._running_lock:
            removed = self._running.pop(task_id, None)
        if removed is not None:
            LOGGER.info("Task removed from running registry | task_id=%s", task_id)