holy-python/ocr_dashboard_service.py

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

import requests

from config import Config
from database import connection
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 list_collections_by_subject(subject_name: str, start_date: str = "", end_date: str = "") -> list[dict[str, Any]]:
    """按主体名查询该主体下所有批次的采集任务。用于同名主体多批次合并展示。"""
    if not subject_name.strip():
        raise TaskError("OCR_SUBJECT_REQUIRED", "subject_name 不能为空", 422)
    return _get(
        f"/xhs/ocr/batches/subject/{subject_name.strip()}/collections",
        {"start_date": start_date, "end_date": end_date},
    ) or []


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


def _overlay_acceptance_questions(data: dict[str, Any], batch_key: str) -> dict[str, Any]:
    """让 Holy蟹看板立即采用同 batch_key 的验收项目当前问题配置。"""
    normalized_key = str(batch_key or "").strip()
    if not normalized_key:
        return data
    try:
        with connection() as conn:
            task = conn.execute(
                """SELECT id,name FROM acceptance_tasks
                   WHERE dashboard_batch_key=? AND scheduled_from_task_id IS NULL
                   ORDER BY id LIMIT 1""",
                (normalized_key,),
            ).fetchone()
            if not task:
                return data
            keywords = [
                row["keyword"] for row in conn.execute(
                    "SELECT keyword FROM acceptance_keywords WHERE task_id=? ORDER BY keyword_order",
                    (task["id"],),
                ).fetchall()
            ]
    except sqlite3.OperationalError:
        return data
    if not keywords:
        return data
    configured_codes: dict[str, str] = {}
    if Config.EVIDENCE_ADMIN_BASE_URL:
        try:
            response = SESSION.get(
                f"{Config.EVIDENCE_ADMIN_BASE_URL}/api/options",
                timeout=min(Config.CRAWLER_TIMEOUT_SECONDS, 10),
            )
            response.raise_for_status()
            scopes = (response.json().get("data") or {}).get("scopes") or []
            configured_codes = {
                str(row.get("search_keyword") or ""): str(row.get("question_code") or "")
                for row in scopes if str(row.get("batch_key") or "") == normalized_key
            }
        except (requests.RequestException, ValueError):
            configured_codes = {}
    current = {
        str(item.get("search_keyword") or ""): item
        for item in data.get("questions", [])
    }
    questions = []
    for index, keyword in enumerate(keywords):
        item = dict(current.get(keyword, {}))
        item.update({
            "question_code": item.get("question_code") or configured_codes.get(keyword)
                             or "X" + hashlib.sha256(keyword.encode()).hexdigest()[:15],
            "search_keyword": keyword,
            "question_direction": item.get("question_direction") or task["name"],
            "display_order": index + 1,
        })
        for field in ("result_count", "device_count", "valid_result_count", "minsheng_matched_count", "cobranded_matched_count"):
            item[field] = int(item.get(field) or 0)
        item.setdefault("minsheng_match_rate", 0)
        item.setdefault("cobranded_match_rate", 0)
        item.setdefault("match_keywords", [])
        questions.append(item)
    data["questions"] = questions
    data["question_pass_count"] = sum(
        1 for item in questions if float(item.get("minsheng_match_rate") or 0) >= 0.5
    )
    return data


def get_series(batch_key: str = "", subject_name: str = "", start_date: str = "", end_date: str = "") -> dict[str, Any]:
    return _get("/xhs/ocr/dashboard/series", {
        "batch_key": batch_key, "subject_name": subject_name, "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"],
                    "collection_key": query.get("collection_key", ""),
                    "question_code": query["question_code"],
                },
                timeout=min(Config.CRAWLER_TIMEOUT_SECONDS, 10),
            )
            response.raise_for_status()
            manual_items = (response.json().get("data") or [])
            data["items"] = manual_items
            data["total"] = len(manual_items)
            data["limit"] = limit
            data["is_complete"] = True
        except (requests.RequestException, ValueError):
            pass
    return data


def list_configured_notes(params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
    """读取凭证后台中由用户配置的多篇展示笔记。"""
    if not Config.EVIDENCE_ADMIN_BASE_URL:
        return []
    try:
        response = SESSION.get(
            f"{Config.EVIDENCE_ADMIN_BASE_URL}/api/notes",
            params={key: value for key, value in (params or {}).items()
                    if key in {"batch_key", "collection_key", "question_code"} and value},
            timeout=min(Config.CRAWLER_TIMEOUT_SECONDS, 15),
        )
        response.raise_for_status()
        payload = response.json()
    except (requests.RequestException, ValueError) as exc:
        raise TaskError("OCR_CONFIGURED_NOTES_UNAVAILABLE", "配置笔记服务暂时不可用", 502) from exc
    if payload.get("error"):
        raise TaskError("OCR_CONFIGURED_NOTES_ERROR", str(payload["error"].get("message") or "配置笔记读取失败"), 502)
    return payload.get("data") or []


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)