ShiErFeng/shierfeng-fastapi-backend/module_business/service/pallet_service.py

import secrets
from datetime import datetime
from io import BytesIO
from typing import Literal, NoReturn

import qrcode
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from common.vo import CrudResponseModel, PageModel
from exceptions.exception import ServiceException, ServiceWarning
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_business.dao.access_dao import AccessDao
from module_business.dao.master_dao import MasterDao
from module_business.domain.test_batch_scope import TestBatchScope, assert_record_in_scope
from module_business.entity.do.master_do import BizPallet
from module_business.entity.do.operation_do import BizPalletInventory
from module_business.entity.vo.master_vo import (
    BusinessUserOptionModel,
    DeletePalletModel,
    PalletCreateModel,
    PalletEditModel,
    PalletModel,
    PalletPageQueryModel,
    UserAssignModel,
)
from module_business.service.test_batch_access_service import TestBatchAccessService

AssignmentResource = Literal['warehouse', 'site']

_ASSIGNMENT_LABELS: dict[AssignmentResource, str] = {
    'warehouse': '仓库',
    'site': '客户现场',
}


class PalletService:
    """托盘资产及业务用户分配服务层。"""

    @staticmethod
    def render_qr_png(qr_token: str) -> bytes:
        """将固定二维码标识渲染为 PNG；补打不会生成新标识。"""
        buffer = BytesIO()
        qrcode.make(qr_token).save(buffer, format='PNG')
        return buffer.getvalue()

    @staticmethod
    async def _prepare_test_write(
        db: AsyncSession,
        current_user: CurrentUserModel,
        scope: TestBatchScope,
    ) -> None:
        if scope.mode != 'TEST':
            return
        TestBatchAccessService.assert_super_admin(current_user)
        await TestBatchAccessService.lock_active_batch(db, scope)

    @classmethod
    async def get_pallet_qrcode(cls, db: AsyncSession, pallet_id: int, scope: TestBatchScope) -> bytes:
        pallet = await MasterDao.get_pallet(db, pallet_id, scope)
        if pallet is None:
            raise ServiceException(message='托盘不存在')
        assert_record_in_scope(pallet, scope, '托盘')
        return cls.render_qr_png(pallet.qr_token)

    @classmethod
    async def _lock_pallet_inventory(
        cls,
        db: AsyncSession,
        pallet_id: int,
        scope: TestBatchScope,
    ) -> tuple[BizPallet, BizPalletInventory]:
        """按库存、托盘顺序加锁，避免与流转事务形成反序死锁。"""
        inventory = await MasterDao.get_pallet_inventory(db, pallet_id, scope, for_update=True)
        if inventory is None:
            if await MasterDao.get_pallet(db, pallet_id, scope) is None:
                raise ServiceException(message='托盘不存在')
            raise ServiceException(message='托盘库存不存在')
        assert_record_in_scope(inventory, scope, '托盘库存')
        pallet = await MasterDao.get_pallet(db, pallet_id, scope, for_update=True)
        if pallet is None:
            raise ServiceException(message='托盘不存在')
        assert_record_in_scope(pallet, scope, '托盘')
        return pallet, inventory

    @classmethod
    async def list_pallet(
        cls,
        db: AsyncSession,
        query_object: PalletPageQueryModel,
        scope: TestBatchScope,
        is_page: bool = True,
    ) -> PageModel | list[PalletModel]:
        result = await MasterDao.list_pallet(db, query_object, scope, is_page)
        records = result.rows if isinstance(result, PageModel) else result
        rows = [PalletModel.model_validate(record) for record in records]
        if isinstance(result, PageModel):
            return result.model_copy(update={'rows': rows})
        return rows

    @classmethod
    async def detail_pallet(cls, db: AsyncSession, pallet_id: int, scope: TestBatchScope) -> PalletModel:
        pallet = await MasterDao.get_pallet(db, pallet_id, scope)
        if pallet is None:
            raise ServiceException(message='托盘不存在')
        assert_record_in_scope(pallet, scope, '托盘')
        return PalletModel.model_validate(pallet)

    @classmethod
    async def add_pallet(
        cls,
        db: AsyncSession,
        pallet: PalletCreateModel,
        current_user: CurrentUserModel,
        scope: TestBatchScope,
    ) -> CrudResponseModel:
        try:
            await cls._prepare_test_write(db, current_user, scope)
            if await MasterDao.find_pallet_by_code(db, pallet.pallet_code) is not None:
                raise ServiceException(message='托盘编号已存在')

            current_time = datetime.now()
            current_user_name = current_user.user.user_name
            pallet_record = BizPallet(
                test_batch_id=scope.batch_id,
                pallet_code=pallet.pallet_code,
                qr_token=secrets.token_urlsafe(32),
                capacity_packages=pallet.capacity_packages,
                asset_status='0',
                remark=pallet.remark,
                create_by=current_user_name,
                create_time=current_time,
                update_by=current_user_name,
                update_time=current_time,
            )
            pallet_record = await MasterDao.add_pallet(db, pallet_record)
            inventory = BizPalletInventory(
                pallet_id=pallet_record.pallet_id,
                test_batch_id=scope.batch_id,
                flow_state='UNSTOCKED',
                remaining_packages=0,
                version=0,
                updated_at=current_time,
            )
            await MasterDao.add_pallet_inventory(db, inventory)
            await db.commit()
            return CrudResponseModel(is_success=True, message='新增成功')
        except IntegrityError:
            await db.rollback()
            raise ServiceException(message='新增托盘失败，托盘编号或二维码标识已存在') from None
        except Exception:
            await db.rollback()
            raise

    @classmethod
    async def edit_pallet(
        cls,
        db: AsyncSession,
        pallet: PalletEditModel,
        current_user: CurrentUserModel,
        scope: TestBatchScope,
    ) -> CrudResponseModel:
        try:
            await cls._prepare_test_write(db, current_user, scope)
            _, inventory = await cls._lock_pallet_inventory(db, pallet.pallet_id, scope)
            if pallet.capacity_packages < inventory.remaining_packages:
                raise ServiceException(message='托盘容量不能小于当前剩余包数')
            await MasterDao.update_pallet(
                db,
                pallet,
                scope,
                update_by=current_user.user.user_name,
                update_time=datetime.now(),
            )
            await db.commit()
            return CrudResponseModel(is_success=True, message='更新成功')
        except IntegrityError:
            await db.rollback()
            raise ServiceException(message='修改托盘失败，托盘状态已变化') from None
        except Exception:
            await db.rollback()
            raise

    @classmethod
    async def remove_pallet(
        cls,
        db: AsyncSession,
        value: int | DeletePalletModel,
    ) -> NoReturn:
        raise ServiceWarning('', '不支持单条物理删除，请通过测试数据批次清理')

    @classmethod
    async def _list_resource_users(
        cls,
        db: AsyncSession,
        resource: AssignmentResource,
        resource_id: int,
        scope: TestBatchScope,
    ) -> list[BusinessUserOptionModel]:
        record = await getattr(MasterDao, f'get_{resource}')(db, resource_id, scope)
        if record is None:
            raise ServiceException(message=f'{_ASSIGNMENT_LABELS[resource]}不存在')
        assert_record_in_scope(record, scope, _ASSIGNMENT_LABELS[resource])

        users = await MasterDao.list_assignable_users(db)
        assigned_user_ids = set(await getattr(MasterDao, f'list_{resource}_assigned_user_ids')(db, resource_id, scope))
        return [
            BusinessUserOptionModel(
                user_id=user.user_id,
                user_name=user.user_name,
                nick_name=user.nick_name,
                assigned=user.user_id in assigned_user_ids,
            )
            for user in users
        ]

    @classmethod
    async def list_warehouse_users(
        cls,
        db: AsyncSession,
        warehouse_id: int,
        scope: TestBatchScope,
    ) -> list[BusinessUserOptionModel]:
        return await cls._list_resource_users(db, 'warehouse', warehouse_id, scope)

    @classmethod
    async def list_site_users(
        cls,
        db: AsyncSession,
        site_id: int,
        scope: TestBatchScope,
    ) -> list[BusinessUserOptionModel]:
        return await cls._list_resource_users(db, 'site', site_id, scope)

    @classmethod
    async def _assign_resource_users(
        cls,
        db: AsyncSession,
        resource: AssignmentResource,
        resource_id: int,
        assignment: UserAssignModel,
        current_user: CurrentUserModel,
        scope: TestBatchScope,
    ) -> CrudResponseModel:
        label = _ASSIGNMENT_LABELS[resource]
        try:
            await cls._prepare_test_write(db, current_user, scope)
            record = await getattr(MasterDao, f'get_{resource}')(
                db,
                resource_id,
                scope,
                for_update=True,
            )
            if record is None:
                raise ServiceException(message=f'{label}不存在')
            assert_record_in_scope(record, scope, label)
            users = await MasterDao.list_assignable_users_by_ids(db, assignment.user_ids)
            if {user.user_id for user in users} != set(assignment.user_ids):
                raise ServiceException(message='存在不存在、停用或已删除的用户')
            await getattr(AccessDao, f'replace_{resource}_users')(db, resource_id, assignment.user_ids)
            await db.commit()
            return CrudResponseModel(is_success=True, message='分配成功')
        except IntegrityError:
            await db.rollback()
            raise ServiceException(message=f'分配{label}用户失败，关联数据无效') from None
        except Exception:
            await db.rollback()
            raise

    @classmethod
    async def assign_warehouse_users(
        cls,
        db: AsyncSession,
        warehouse_id: int,
        assignment: UserAssignModel,
        current_user: CurrentUserModel,
        scope: TestBatchScope,
    ) -> CrudResponseModel:
        return await cls._assign_resource_users(
            db,
            'warehouse',
            warehouse_id,
            assignment,
            current_user,
            scope,
        )

    @classmethod
    async def assign_site_users(
        cls,
        db: AsyncSession,
        site_id: int,
        assignment: UserAssignModel,
        current_user: CurrentUserModel,
        scope: TestBatchScope,
    ) -> CrudResponseModel:
        return await cls._assign_resource_users(
            db,
            'site',
            site_id,
            assignment,
            current_user,
            scope,
        )