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

"""运输管理 - 侧栏菜单角标统计"""
from __future__ import annotations

from collections import defaultdict
from typing import Dict, List, Optional

from sqlalchemy.orm import Session

from app.models.pl_models import (
    PlDocumentReminder,
    PlFreightSettlement,
    PlPlan,
    PlPlanProduct,
    PlVehicleMaintenanceApply,
    PlVehicleRechargeApply,
    PlVehicleUreaApply,
)
from app.services.pl_cert_service import scan_cert_expiry
from app.services.pl_data_scope import apply_pl_scope
from app.services.pl_freight_settlement_service import can_user_approve_settlement
from app.services.pl_plan_service import plan_has_hazard_products

MENU_CODE_DOCUMENT = "pl:document:list"
MENU_CODE_FREIGHT_APPROVAL = "pl:freight-approval:list"
MENU_CODE_FREIGHT_DISBURSE = "pl:freight-disburse:list"
MENU_CODE_RECHARGE_RECORD = "pl:vehicle-recharge-record:list"
MENU_CODE_MAINT_PROCESS = "pl:vehicle-maintenance-process:list"
MENU_CODE_UREA_DISBURSE = "pl:vehicle-urea-disburse:list"
MENU_CODE_PLAN_CERT = "pl:plan-cert:list"
MENU_CODE_PLAN_ARCHIVE = "pl:plan-archive:list"

SUPPORTED_MENU_CODES = {
    MENU_CODE_DOCUMENT,
    MENU_CODE_FREIGHT_APPROVAL,
    MENU_CODE_FREIGHT_DISBURSE,
    MENU_CODE_RECHARGE_RECORD,
    MENU_CODE_MAINT_PROCESS,
    MENU_CODE_UREA_DISBURSE,
    MENU_CODE_PLAN_CERT,
    MENU_CODE_PLAN_ARCHIVE,
}


def _count_pending_applies(db: Session, tenant_no: str, user, model) -> int:
    query = db.query(model).filter(
        model.tenant_no == tenant_no,
        model.del_flag == 0,
        model.status == "pending",
    )
    query = apply_pl_scope(query, model, user, db, tenant_no)
    return query.count()


def count_document_reminder_expired_pending(db: Session, tenant_no: str, user) -> int:
    scan_cert_expiry(db, tenant_no)
    db.flush()
    query = db.query(PlDocumentReminder).filter(
        PlDocumentReminder.tenant_no == tenant_no,
        PlDocumentReminder.del_flag == 0,
        PlDocumentReminder.status == "pending",
        PlDocumentReminder.days_remaining <= 0,
    )
    query = apply_pl_scope(query, PlDocumentReminder, user, db, tenant_no)
    return query.count()


def count_freight_approval_pending(db: Session, tenant_no: str, user_id: str) -> int:
    if not user_id:
        return 0
    rows = (
        db.query(PlFreightSettlement.id)
        .filter(
            PlFreightSettlement.tenant_no == tenant_no,
            PlFreightSettlement.del_flag == 0,
            PlFreightSettlement.settlement_status == "approving",
        )
        .all()
    )
    total = 0
    for (settlement_id,) in rows:
        if can_user_approve_settlement(db, settlement_id, user_id):
            total += 1
    return total


def count_plan_cert_pending(db: Session, tenant_no: str, user) -> int:
    query = db.query(PlPlan).filter(
        PlPlan.tenant_no == tenant_no,
        PlPlan.del_flag == 0,
        PlPlan.hazard_cert_ready == 0,
        PlPlan.archive_status != "archived",
        PlPlan.follow_status != "cancelled",
    )
    query = apply_pl_scope(query, PlPlan, user, db, tenant_no)
    plans = query.all()
    if not plans:
        return 0

    plan_ids = [plan.id for plan in plans]
    products_by_plan: dict[str, list[PlPlanProduct]] = defaultdict(list)
    product_rows = (
        db.query(PlPlanProduct)
        .filter(
            PlPlanProduct.plan_id.in_(plan_ids),
            PlPlanProduct.del_flag == 0,
        )
        .all()
    )
    for product in product_rows:
        products_by_plan[str(product.plan_id)].append(product)

    total = 0
    for plan in plans:
        products = products_by_plan.get(str(plan.id), [])
        if plan_has_hazard_products(db, tenant_no, products):
            total += 1
    return total


def count_plan_archive_pending(db: Session, tenant_no: str, user) -> int:
    query = db.query(PlPlan).filter(
        PlPlan.tenant_no == tenant_no,
        PlPlan.del_flag == 0,
        PlPlan.archive_status == "unarchived",
        PlPlan.follow_status != "cancelled",
    )
    query = apply_pl_scope(query, PlPlan, user, db, tenant_no)
    return query.count()


def count_freight_disburse_pending(db: Session, tenant_no: str, user) -> int:
    query = db.query(PlFreightSettlement).filter(
        PlFreightSettlement.tenant_no == tenant_no,
        PlFreightSettlement.del_flag == 0,
        PlFreightSettlement.settlement_status == "approved",
    )
    query = apply_pl_scope(query, PlFreightSettlement, user, db, tenant_no)
    return query.count()


def build_menu_badge_counts(
    db: Session,
    *,
    tenant_no: str,
    user,
    user_id: str,
    menu_codes: Optional[List[str]] = None,
) -> Dict[str, int]:
    requested = {c.strip() for c in (menu_codes or []) if c and c.strip()}
    if not requested:
        requested = set(SUPPORTED_MENU_CODES)
    requested &= SUPPORTED_MENU_CODES

    counts: Dict[str, int] = {}
    if MENU_CODE_DOCUMENT in requested:
        counts[MENU_CODE_DOCUMENT] = count_document_reminder_expired_pending(db, tenant_no, user)
    if MENU_CODE_FREIGHT_APPROVAL in requested:
        counts[MENU_CODE_FREIGHT_APPROVAL] = count_freight_approval_pending(db, tenant_no, user_id)
    if MENU_CODE_FREIGHT_DISBURSE in requested:
        counts[MENU_CODE_FREIGHT_DISBURSE] = count_freight_disburse_pending(db, tenant_no, user)
    if MENU_CODE_RECHARGE_RECORD in requested:
        counts[MENU_CODE_RECHARGE_RECORD] = _count_pending_applies(
            db, tenant_no, user, PlVehicleRechargeApply
        )
    if MENU_CODE_MAINT_PROCESS in requested:
        counts[MENU_CODE_MAINT_PROCESS] = _count_pending_applies(
            db, tenant_no, user, PlVehicleMaintenanceApply
        )
    if MENU_CODE_UREA_DISBURSE in requested:
        counts[MENU_CODE_UREA_DISBURSE] = _count_pending_applies(
            db, tenant_no, user, PlVehicleUreaApply
        )
    if MENU_CODE_PLAN_CERT in requested:
        counts[MENU_CODE_PLAN_CERT] = count_plan_cert_pending(db, tenant_no, user)
    if MENU_CODE_PLAN_ARCHIVE in requested:
        counts[MENU_CODE_PLAN_ARCHIVE] = count_plan_archive_pending(db, tenant_no, user)
    return counts