shanxia2.0/backend/app/services/pl_ocr_service.py

"""阿里云人员证书 OCR（身份证 / 驾驶证）"""
from __future__ import annotations

import io
import json
import logging
import re
from datetime import datetime
from typing import Any, Optional

from fastapi import HTTPException
from sqlalchemy.orm import Session

from app.core.config import settings
from app.core.utils import generate_uuid
from app.models.pl_models import PlOcrUsageMonth

logger = logging.getLogger(__name__)

SUPPORTED_CERT_NAMES = ("身份证", "驾驶证", "行驶证")
MAX_UPLOAD_BYTES = 1_500_000


def _monthly_limit() -> int:
    return max(1, int(settings.ALIYUN_OCR_MONTHLY_LIMIT or 170))


def current_usage_month() -> str:
    return datetime.now().strftime("%Y-%m")


def get_or_create_usage_row(db: Session, usage_month: str) -> PlOcrUsageMonth:
    row = (
        db.query(PlOcrUsageMonth)
        .filter(PlOcrUsageMonth.usage_month == usage_month, PlOcrUsageMonth.del_flag == 0)
        .first()
    )
    if row:
        return row
    row = PlOcrUsageMonth(
        id=generate_uuid(),
        usage_month=usage_month,
        used_count=0,
        tenant_no="shanxia2.0",
        create_dept_id="",
        del_flag=0,
    )
    db.add(row)
    db.flush()
    return row


def assert_quota_available(db: Session) -> PlOcrUsageMonth:
    month = current_usage_month()
    row = get_or_create_usage_row(db, month)
    limit = _monthly_limit()
    if (row.used_count or 0) >= limit:
        raise HTTPException(
            status_code=400,
            detail=f"本月 OCR 识别次数已达上限（{limit} 次），请下月再试或联系管理员",
        )
    return row


def increment_usage(db: Session, row: PlOcrUsageMonth) -> int:
    row.used_count = int(row.used_count or 0) + 1
    row.update_time = datetime.now()
    db.add(row)
    db.commit()
    db.refresh(row)
    return int(row.used_count or 0)


def _first_attachment_url(attachment_urls: str) -> str:
    parts = [p.strip() for p in (attachment_urls or "").split(",") if p.strip()]
    if not parts:
        raise HTTPException(status_code=400, detail="请先上传证书附件")
    url = parts[0]
    lower = url.lower()
    if lower.endswith(".pdf"):
        raise HTTPException(status_code=400, detail="OCR 暂不支持 PDF，请上传图片附件")
    return url


def compress_image_bytes(raw: bytes) -> bytes:
    if len(raw) <= MAX_UPLOAD_BYTES:
        return raw
    from PIL import Image

    img = Image.open(io.BytesIO(raw))
    if img.mode not in ("RGB", "L"):
        img = img.convert("RGB")
    quality = 85
    width, height = img.size
    buf = io.BytesIO()
    while True:
        buf.seek(0)
        buf.truncate(0)
        img.save(buf, format="JPEG", quality=quality, optimize=True)
        size = buf.tell()
        if size <= MAX_UPLOAD_BYTES or quality <= 40:
            return buf.getvalue()
        quality -= 10
        if size > MAX_UPLOAD_BYTES * 1.5:
            width = int(width * 0.85)
            height = int(height * 0.85)
            img = img.resize((width, height), Image.Resampling.LANCZOS)


def parse_expire_date(valid_period: Optional[str]) -> Optional[str]:
    if not valid_period:
        return None
    text = str(valid_period).strip()
    if "长期" in text:
        return None
    m = re.search(
        r"(\d{4})[.\-/年](\d{1,2})[.\-/月](\d{1,2}).*?[至\-—~]\s*(\d{4})[.\-/年](\d{1,2})[.\-/月](\d{1,2})",
        text,
    )
    if m:
        y, mo, d = m.group(4), m.group(5), m.group(6)
        return f"{int(y):04d}-{int(mo):02d}-{int(d):02d}"
    m2 = re.search(r"(\d{4})[.\-/](\d{1,2})[.\-/](\d{1,2})\s*$", text)
    if m2:
        return f"{int(m2.group(1)):04d}-{int(m2.group(2)):02d}-{int(m2.group(3)):02d}"
    return None


def parse_vehicle_license_expire(inspection_record: Optional[str]) -> Optional[str]:
    """从行驶证副页检验记录解析有效期结束日。

    常见文案：检验有效期至2026年01月 / 至2026-01 / 2026年1月
    仅到年月时默认取当月最后一天的近似：用该月 01 日写入（由运营可手改），
    或优先匹配完整年月日。
    """
    if not inspection_record:
        return None
    text = str(inspection_record).strip()
    # 完整日期
    full = parse_expire_date(text)
    if full:
        return full
    # 检验有效期至YYYY年MM月 / 至YYYY-MM
    m = re.search(
        r"(?:有效期|检验).*?至\s*(\d{4})[.\-/年](\d{1,2})月?",
        text,
    )
    if not m:
        m = re.search(r"至\s*(\d{4})[.\-/年](\d{1,2})月?", text)
    if m:
        y, mo = int(m.group(1)), int(m.group(2))
        return f"{y:04d}-{mo:02d}-01"
    return None


def _build_ocr_client():
    ak = (settings.ALIYUN_OCR_ACCESS_KEY_ID or "").strip()
    sk = (settings.ALIYUN_OCR_ACCESS_KEY_SECRET or "").strip()
    endpoint = (settings.ALIYUN_OCR_ENDPOINT or "ocr-api.cn-hangzhou.aliyuncs.com").strip()
    if not ak or not sk:
        raise HTTPException(status_code=400, detail="OCR 未配置，请联系管理员")
    try:
        from alibabacloud_ocr_api20210707.client import Client as OcrClient
        from alibabacloud_tea_openapi import models as open_api_models
    except ImportError as exc:
        logger.exception("阿里云 OCR SDK 未安装")
        raise HTTPException(status_code=500, detail="OCR 组件未安装，请联系管理员") from exc
    return OcrClient(
        open_api_models.Config(
            access_key_id=ak,
            access_key_secret=sk,
            endpoint=endpoint,
        )
    )


def _parse_aliyun_data(raw_data: Any) -> dict:
    if isinstance(raw_data, str):
        return json.loads(raw_data) if raw_data else {}
    if isinstance(raw_data, dict):
        return raw_data
    return {}


def _recognize_idcard(image_bytes: bytes) -> dict:
    from alibabacloud_ocr_api20210707 import models as ocr_models
    from alibabacloud_tea_util import models as util_models

    client = _build_ocr_client()
    request = ocr_models.RecognizeIdcardRequest(body=io.BytesIO(image_bytes))
    resp = client.recognize_idcard_with_options(request, util_models.RuntimeOptions())
    body = resp.body
    if getattr(body, "code", None):
        raise HTTPException(
            status_code=400,
            detail=getattr(body, "message", None) or "身份证识别失败",
        )
    data = _parse_aliyun_data(body.data)
    face = (((data or {}).get("data") or {}).get("face") or {}).get("data") or {}
    back = (((data or {}).get("data") or {}).get("back") or {}).get("data") or {}
    cert_no = face.get("idNumber") or back.get("idNumber")
    valid_period = back.get("validPeriod") or face.get("validPeriod")
    expire_date = parse_expire_date(valid_period)
    if not cert_no and not expire_date:
        raise HTTPException(status_code=400, detail="未能识别到身份证证号或有效期，请换清晰图片重试")
    return {
        "certNo": str(cert_no).strip() if cert_no else None,
        "expireDate": expire_date,
        "expireDateRaw": valid_period,
    }


def _recognize_driving_license(image_bytes: bytes) -> dict:
    from alibabacloud_ocr_api20210707 import models as ocr_models
    from alibabacloud_tea_util import models as util_models

    client = _build_ocr_client()
    request = ocr_models.RecognizeDrivingLicenseRequest(body=io.BytesIO(image_bytes))
    resp = client.recognize_driving_license_with_options(request, util_models.RuntimeOptions())
    body = resp.body
    if getattr(body, "code", None):
        raise HTTPException(
            status_code=400,
            detail=getattr(body, "message", None) or "驾驶证识别失败",
        )
    data = _parse_aliyun_data(body.data)
    face = (((data or {}).get("data") or {}).get("face") or {}).get("data") or {}
    cert_no = face.get("licenseNumber")
    valid_period = face.get("validPeriod")
    expire_date = parse_expire_date(valid_period)
    if not cert_no and not expire_date:
        raise HTTPException(status_code=400, detail="未能识别到驾驶证证号或有效期，请换清晰图片重试")
    return {
        "certNo": str(cert_no).strip() if cert_no else None,
        "expireDate": expire_date,
        "expireDateRaw": valid_period,
    }


def _recognize_vehicle_license(image_bytes: bytes) -> dict:
    """行驶证：证号优先车辆识别代号(VIN)，有效期从副页检验记录解析。"""
    from alibabacloud_ocr_api20210707 import models as ocr_models
    from alibabacloud_tea_util import models as util_models

    client = _build_ocr_client()
    request = ocr_models.RecognizeVehicleLicenseRequest(body=io.BytesIO(image_bytes))
    resp = client.recognize_vehicle_license_with_options(request, util_models.RuntimeOptions())
    body = resp.body
    if getattr(body, "code", None):
        raise HTTPException(
            status_code=400,
            detail=getattr(body, "message", None) or "行驶证识别失败",
        )
    data = _parse_aliyun_data(body.data)
    face = (((data or {}).get("data") or {}).get("face") or {}).get("data") or {}
    back = (((data or {}).get("data") or {}).get("back") or {}).get("data") or {}
    cert_no = face.get("vinCode") or back.get("vinCode") or face.get("licensePlateNumber")
    inspection_record = back.get("inspectionRecord") or face.get("inspectionRecord")
    expire_date = parse_vehicle_license_expire(inspection_record)
    if not cert_no and not expire_date:
        raise HTTPException(status_code=400, detail="未能识别到行驶证证号或有效期，请换清晰图片重试")
    return {
        "certNo": str(cert_no).strip() if cert_no else None,
        "expireDate": expire_date,
        "expireDateRaw": inspection_record,
    }


def recognize_personnel_cert(
    db: Session,
    *,
    cert_name: str,
    attachment_urls: str,
) -> dict:
    name = (cert_name or "").strip()
    if name not in SUPPORTED_CERT_NAMES:
        raise HTTPException(status_code=400, detail="仅支持身份证、驾驶证、行驶证自动识别")

    url = _first_attachment_url(attachment_urls)
    try:
        from app.services.bx.storage import download_file_bytes

        raw = download_file_bytes(url)
    except Exception as exc:
        logger.exception("下载 OCR 附件失败: %s", url)
        raise HTTPException(status_code=400, detail="附件读取失败，请重新上传后再试") from exc

    image_bytes = compress_image_bytes(raw)
    usage_row = assert_quota_available(db)
    # 只要准备调用阿里云即扣次（含识别失败），避免免费额度被刷完却不记账
    used = increment_usage(db, usage_row)

    try:
        if name == "身份证":
            result = _recognize_idcard(image_bytes)
        elif name == "驾驶证":
            result = _recognize_driving_license(image_bytes)
        else:
            result = _recognize_vehicle_license(image_bytes)
    except HTTPException:
        raise
    except Exception as exc:
        logger.exception("阿里云 OCR 调用失败 cert_name=%s", name)
        raise HTTPException(status_code=400, detail="OCR 识别失败，请稍后重试") from exc

    result["monthlyUsed"] = used
    result["monthlyLimit"] = _monthly_limit()
    return result