from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from module_business.entity.do.master_do import BizPallet
from module_business.entity.do.operation_do import BizPalletEvent, BizPalletInventory


class InventoryDao:
    """托盘库存锁、资产锁和只追加事件数据库操作层。"""

    @classmethod
    async def get_inventory_for_update(
        cls,
        db: AsyncSession,
        pallet_id: int,
    ) -> BizPalletInventory | None:
        result = await db.execute(
            select(BizPalletInventory).where(BizPalletInventory.pallet_id == pallet_id).with_for_update()
        )
        return result.scalar_one_or_none()

    @classmethod
    async def get_inventories_for_update(
        cls,
        db: AsyncSession,
        pallet_ids: list[int],
    ) -> list[BizPalletInventory]:
        ordered_ids = sorted(set(pallet_ids))
        if not ordered_ids:
            return []

        result = await db.execute(
            select(BizPalletInventory)
            .where(BizPalletInventory.pallet_id.in_(ordered_ids))
            .order_by(BizPalletInventory.pallet_id)
            .with_for_update()
        )
        return sorted(result.scalars().all(), key=lambda inventory: inventory.pallet_id)

    @classmethod
    async def get_pallet_by_qr_token(
        cls,
        db: AsyncSession,
        qr_token: str,
    ) -> BizPallet | None:
        result = await db.execute(select(BizPallet).where(BizPallet.qr_token == qr_token))
        return result.scalar_one_or_none()

    @classmethod
    async def get_pallets_by_ids(
        cls,
        db: AsyncSession,
        pallet_ids: list[int],
    ) -> list[BizPallet]:
        """在调用方已锁库存后，按托盘 ID 升序锁定资产行。"""
        ordered_ids = sorted(set(pallet_ids))
        if not ordered_ids:
            return []

        result = await db.execute(
            select(BizPallet)
            .where(BizPallet.pallet_id.in_(ordered_ids))
            .order_by(BizPallet.pallet_id)
            .with_for_update()
        )
        return sorted(result.scalars().all(), key=lambda pallet: pallet.pallet_id)

    @classmethod
    async def get_event_by_idempotency_key(
        cls,
        db: AsyncSession,
        pallet_id: int,
        idempotency_key: str,
    ) -> BizPalletEvent | None:
        result = await db.execute(
            select(BizPalletEvent).where(
                BizPalletEvent.pallet_id == pallet_id,
                BizPalletEvent.idempotency_key == idempotency_key,
            )
        )
        return result.scalar_one_or_none()

    @classmethod
    async def add_event(
        cls,
        db: AsyncSession,
        inventory: BizPalletInventory,
        event: BizPalletEvent,
    ) -> BizPalletEvent:
        """在调用方持有库存行锁时，为托盘追加下一序号事件。"""
        sequence_result = await db.execute(
            select(func.max(BizPalletEvent.event_seq)).where(BizPalletEvent.pallet_id == inventory.pallet_id)
        )
        last_sequence = sequence_result.scalar_one()
        event.pallet_id = inventory.pallet_id
        event.event_seq = (last_sequence or 0) + 1
        db.add(event)
        await db.flush()
        inventory.last_event_id = event.event_id
        await db.flush()
        return event