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 = "",
                 start_date: str = "", end_date: str = "") -> dict[str, Any]:
    return _get("/xhs/ocr/dashboard/summary", {
        "batch_key": batch_key, "collection_key": collection_key,
        "start_date": start_date, "end_date": end_date,
    }) or {}


def get_series(batch_key: str = "", start_date: str = "", end_date: str = "") -> dict[str, Any]:
    return _get("/xhs/ocr/dashboard/series", {
        "batch_key": batch_key, "start_date": start_date, "end_date": end_date,
    }) or {"items": [], "subjects": []}


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


def list_results(params: dict[str, Any]) -> dict[str, Any]:
    allowed = {
        "batch_key", "collection_key", "start_date", "end_date",
        "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", "start_date", "end_date", "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 str(query.get("collection_key") or "").strip()
            or str(query.get("start_date") or "").strip() or str(query.get("end_date") or "").strip()):
        raise TaskError("OCR_EVIDENCE_SCOPE_REQUIRED", "batch_key/collection_key 或 start_date/end_date 至少传入一种", 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
    data = _get("/xhs/ocr/dashboard/evidences", query) or {
        "items": [], "total": 0, "limit": limit, "is_complete": False,
    }
    if Config.EVIDENCE_ADMIN_BASE_URL and query.get("batch_key") and query.get("question_code"):
        try:
            response = SESSION.get(
                f"{Config.EVIDENCE_ADMIN_BASE_URL}/api/evidences",
                params={"batch_key": query["batch_key"], "question_code": query["question_code"]},
                timeout=min(Config.CRAWLER_TIMEOUT_SECONDS, 10),
            )
            response.raise_for_status()
            manual_items = (response.json().get("data") or [])
            items = data.setdefault("items", [])
            existing_urls = {str(item.get("cos_url") or "") for item in items}
            object_indexes = {
                str(item.get("cos_object_key")): index for index, item in enumerate(items)
                if item.get("cos_object_key")
            }
            for item in manual_items:
                url = str(item.get("cos_url") or "")
                object_key = str(item.get("cos_object_key") or "")
                if object_key and object_key in object_indexes:
                    items[object_indexes[object_key]] = item
                    existing_urls.add(url)
                elif url and url not in existing_urls:
                    items.append(item)
                    if object_key:
                        object_indexes[object_key] = len(items) - 1
                    existing_urls.add(url)
            evidence_type_order = {
                "body_ranking": 0, "first_screen": 1, "reference_source": 2,
                "before": 3, "after": 3, "before_after": 3, "viral_note": 4,
            }
            items.sort(key=lambda item: (
                evidence_type_order.get(str(item.get("evidence_type") or ""), 99),
                1 if (
                    item.get("image_id") is not None
                    or "/manual-evidences/" in str(item.get("cos_object_key") or "")
                    or item.get("upload_source") == "holy-backend"
                ) else 0,
                str(item.get("created_at") or ""),
                int(item.get("sort_order") or 0),
                int(item.get("image_id") or item.get("evidence_id") or item.get("id") or 0),
            ))
            data["total"] = len(data.get("items") or [])
        except (requests.RequestException, ValueError):
            pass
    return data


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)