certificate-system/backend/app/services/pdf_pregeneration.py

import queue
import threading
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path

from app.db.session import SessionLocal
from app.models import Certificate, CertificateAccessToken, Learner, ProjectCourse
from app.services.pdf import cached_pdf_is_fresh, certificate_template_path, pdf_cache_path, render_certificate_pdf

MAX_PDF_PREGENERATION_JOBS = 30


def _utc_now() -> datetime:
    return datetime.now(UTC).replace(tzinfo=None)


@dataclass
class PdfPregenerationJob:
    id: str
    certificate_ids: list[int]
    concurrency_limit: int
    status: str = "running"
    total: int = 0
    completed: int = 0
    generated: int = 0
    cached: int = 0
    skipped: int = 0
    failed: int = 0
    errors: list[str] = field(default_factory=list)
    created_at: datetime = field(default_factory=_utc_now)
    finished_at: datetime | None = None

    def to_dict(self) -> dict[str, object]:
        percent = round((self.completed / self.total) * 100) if self.total else 0
        return {
            "id": self.id,
            "status": self.status,
            "total": self.total,
            "completed": self.completed,
            "generated": self.generated,
            "cached": self.cached,
            "skipped": self.skipped,
            "failed": self.failed,
            "percent": percent,
            "concurrency_limit": self.concurrency_limit,
            "errors": self.errors[-10:],
            "created_at": self.created_at,
            "finished_at": self.finished_at,
        }


class PdfPregenerationManager:
    def __init__(self) -> None:
        self._jobs: dict[str, PdfPregenerationJob] = {}
        self._lock = threading.Lock()

    def start(self, certificate_ids: list[int], concurrency_limit: int) -> PdfPregenerationJob:
        unique_ids = _unique_ids(certificate_ids)
        if not unique_ids:
            raise ValueError("请选择需要预生成PDF的证书")
        job = PdfPregenerationJob(
            id=uuid.uuid4().hex,
            certificate_ids=unique_ids,
            concurrency_limit=max(1, concurrency_limit),
            total=len(unique_ids),
        )
        with self._lock:
            self._jobs[job.id] = job
            self._trim_jobs()
        thread = threading.Thread(target=self._run_job, args=(job.id,), daemon=True)
        thread.start()
        return job

    def get(self, job_id: str) -> PdfPregenerationJob | None:
        with self._lock:
            return self._jobs.get(job_id)

    def list(self) -> list[PdfPregenerationJob]:
        with self._lock:
            return sorted(self._jobs.values(), key=lambda item: item.created_at, reverse=True)

    def _run_job(self, job_id: str) -> None:
        job = self.get(job_id)
        if not job:
            return

        work_queue: queue.Queue[int] = queue.Queue()
        for certificate_id in job.certificate_ids:
            work_queue.put(certificate_id)

        workers = [
            threading.Thread(target=self._worker, args=(job_id, work_queue), daemon=True)
            for _ in range(min(job.concurrency_limit, job.total))
        ]
        for worker in workers:
            worker.start()
        for worker in workers:
            worker.join()

        with self._lock:
            job.status = "completed_with_errors" if job.failed else "completed"
            job.finished_at = _utc_now()

    def _worker(self, job_id: str, work_queue: queue.Queue[int]) -> None:
        while True:
            try:
                certificate_id = work_queue.get_nowait()
            except queue.Empty:
                return
            try:
                job = self.get(job_id)
                if not job:
                    return
                result = pre_generate_certificate_pdf(certificate_id, job.concurrency_limit)
                self._record(job_id, result)
            except Exception as exc:
                self._record(job_id, "failed", f"证书ID {certificate_id}：{exc}")
            finally:
                work_queue.task_done()

    def _record(self, job_id: str, result: str, error: str | None = None) -> None:
        job = self.get(job_id)
        if not job:
            return
        with self._lock:
            job.completed += 1
            if result == "generated":
                job.generated += 1
            elif result == "cached":
                job.cached += 1
            elif result == "skipped":
                job.skipped += 1
            else:
                job.failed += 1
                if error:
                    job.errors.append(error)

    def _trim_jobs(self) -> None:
        if len(self._jobs) <= MAX_PDF_PREGENERATION_JOBS:
            return
        old_jobs = sorted(self._jobs.values(), key=lambda item: item.created_at)
        for job in old_jobs[: len(self._jobs) - MAX_PDF_PREGENERATION_JOBS]:
            self._jobs.pop(job.id, None)


def pre_generate_certificate_pdf(certificate_id: int, concurrency_limit: int) -> str:
    with SessionLocal() as db:
        certificate = db.get(Certificate, certificate_id)
        if not certificate or certificate.status != "valid":
            return "skipped"
        learner = db.get(Learner, certificate.learner_id)
        token = db.get(CertificateAccessToken, certificate.qr_token_id or certificate.public_token_id)
        project = db.query(ProjectCourse).filter(ProjectCourse.code == certificate.project_code).first()
        if not learner or not token:
            raise ValueError("证书数据不完整")

        was_cached = certificate_pdf_is_ready(certificate)
        render_certificate_pdf(certificate, learner, project, token, concurrency_limit=concurrency_limit)
        db.commit()
        return "cached" if was_cached else "generated"


def certificate_pdf_is_ready(certificate: Certificate) -> bool:
    output_path = pdf_cache_path(certificate.certificate_no)
    template_path = certificate_template_path()
    latest_renderer_mtime = max(template_path.stat().st_mtime, Path(__file__).with_name("pdf.py").stat().st_mtime)
    return cached_pdf_is_fresh(output_path) and output_path.stat().st_mtime >= latest_renderer_mtime


def _unique_ids(values: list[int]) -> list[int]:
    seen: set[int] = set()
    result: list[int] = []
    for value in values:
        if value not in seen:
            seen.add(value)
            result.append(value)
    return result


pdf_pregeneration_manager = PdfPregenerationManager()