holy-python/ocr_dashboard_service.py

import io
import re
import zipfile
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import requests

from config import Config
from task_service import TaskError


SESSION = requests.Session()
SESSION.trust_env = False


def _get(path: str, params: dict[str, Any] | None = None) -> Any:
    try:
        response = SESSION.get(
            f"{Config.CRAWLER_BASE_URL}/api/v1{path}",
            params={key: value for key, value in (params or {}).items() if value not in (None, "")},
            timeout=Config.CRAWLER_TIMEOUT_SECONDS,
        )
        response.raise_for_status()
        payload = response.json()
    except (requests.RequestException, ValueError) as exc:
        raise TaskError("OCR_DASHBOARD_UNAVAILABLE", "OCR 看板数据服务暂时不可用", 502) from exc
    if not payload.get("success") or int(payload.get("code") or 0) != 200:
        code = int(payload.get("code") or 500)
        raise TaskError(
            "OCR_DASHBOARD_UPSTREAM_ERROR",
            str(payload.get("msg") or "OCR 看板数据查询失败"),
            404 if code == 404 else 422 if code == 400 else 502,
        )
    return payload.get("data")


def list_batches() -> list[dict[str, Any]]:
    return _get("/xhs/ocr/batches") or []


def list_collections(batch_key: str) -> list[dict[str, Any]]:
    if not batch_key.strip():
        raise TaskError("OCR_BATCH_REQUIRED", "batch_key 不能为空", 422)
    return _get(f"/xhs/ocr/batches/{batch_key.strip()}/collections") or []


def get_summary(batch_key: str = "", collection_key: str = "") -> dict[str, Any]:
    return _get("/xhs/ocr/dashboard/summary", {
        "batch_key": batch_key, "collection_key": collection_key,
    }) or {}


def get_comparison(batch_key: str = "", collection_key: str = "") -> dict[str, Any]:
    return _get("/xhs/ocr/dashboard/comparison", {
        "batch_key": batch_key, "collection_key": collection_key,
    }) or {}


def list_results(params: dict[str, Any]) -> dict[str, Any]:
    allowed = {
        "batch_key", "collection_key", "question_code", "search_keyword", "city",
        "department_name", "minsheng_matched", "cobranded_matched", "is_valid",
        "page", "page_size",
    }
    query = {key: value for key, value in params.items() if key in allowed}
    page = int(query.get("page") or 1)
    page_size = int(query.get("page_size") or 50)
    if page < 1 or not 1 <= page_size <= 200:
        raise TaskError("OCR_INVALID_PAGINATION", "分页参数无效", 422)
    query.update({"page": page, "page_size": page_size})
    return _get("/xhs/ocr/dashboard/results", query) or {
        "items": [], "total": 0, "page": page, "page_size": page_size,
    }


def get_result_matches(result_id: int) -> dict[str, Any]:
    if result_id < 1:
        raise TaskError("OCR_RESULT_ID_INVALID", "result_id 必须大于 0", 422)
    return _get(f"/xhs/ocr/results/{result_id}/matches") or {}


def list_evidences(params: dict[str, Any]) -> dict[str, Any]:
    allowed = {"batch_key", "collection_key", "question_code", "evidence_type", "limit"}
    query = {key: value for key, value in params.items() if key in allowed}
    if not str(query.get("batch_key") or "").strip() or not str(query.get("collection_key") or "").strip():
        raise TaskError("OCR_EVIDENCE_SCOPE_REQUIRED", "batch_key 和 collection_key 不能为空", 422)
    limit = int(query.get("limit") or 10)
    if not 1 <= limit <= 100:
        raise TaskError("OCR_EVIDENCE_LIMIT_INVALID", "limit 必须在 1 到 100 之间", 422)
    query["limit"] = limit
    return _get("/xhs/ocr/dashboard/evidences", query) or {
        "items": [], "total": 0, "limit": limit, "is_complete": False,
    }


def export_evidences(params: dict[str, Any]) -> tuple[io.BytesIO, str, int]:
    query = dict(params)
    query["limit"] = 100
    data = list_evidences(query)
    items = data.get("items") or []
    if not items:
        raise TaskError("OCR_EVIDENCE_EMPTY", "当前筛选范围没有可导出的截图", 404)
    output = io.BytesIO()
    used_names: set[str] = set()
    with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as archive:
        for index, item in enumerate(items, 1):
            url = str(item.get("cos_url") or "").strip()
            host = (urlparse(url).hostname or "").lower()
            if not url.startswith("https://") or not host.endswith(".myqcloud.com"):
                continue
            try:
                response = SESSION.get(url, timeout=Config.CRAWLER_TIMEOUT_SECONDS)
                response.raise_for_status()
            except requests.RequestException:
                continue
            suffix = Path(str(item.get("source_file_name") or urlparse(url).path)).suffix.lower()
            if suffix not in {".jpg", ".jpeg", ".png", ".webp"}:
                suffix = ".jpg"
            parts = [item.get("question_code"), item.get("search_keyword"),
                     item.get("device_code"), item.get("city"), str(item.get("sort_order") or index)]
            stem = "_".join(re.sub(r"[^0-9A-Za-z\u4e00-\u9fff_-]+", "-", str(part))
                            for part in parts if part not in (None, ""))
            name = f"{stem or f'evidence-{index}'}{suffix}"
            if name in used_names:
                name = f"{Path(name).stem}-{index}{suffix}"
            used_names.add(name)
            archive.writestr(name, response.content)
    if not used_names:
        raise TaskError("OCR_EVIDENCE_EXPORT_FAILED", "截图下载失败，未生成导出文件", 502)
    output.seek(0)
    scope = str(query.get("question_code") or "all")
    return output, f"holy-crab-evidences-{scope}.zip", len(used_names)