from collections.abc import Collection

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

from exceptions.exception import ServiceException
from module_business.domain.test_batch_graph import BATCH_MODELS, BatchSelection
from module_business.entity.do.master_do import BizUserSite, BizUserWarehouse


class TestBatchPurgeDao:
    """测试批次清理范围查询与无提交删除。"""

    @classmethod
    async def load_selection(
        cls,
        db: AsyncSession,
        batch_id: int,
        *,
        for_update: bool,
    ) -> BatchSelection:
        selection: BatchSelection = {}
        for table_name, (model, primary_key, _) in BATCH_MODELS.items():
            statement = select(primary_key).where(model.test_batch_id == batch_id).order_by(primary_key)
            if for_update:
                statement = statement.with_for_update()
            result = await db.execute(statement)
            selection[table_name] = frozenset(row[0] for row in result.all())
        return selection

    @classmethod
    async def list_display_keys(
        cls,
        db: AsyncSession,
        selection: BatchSelection,
    ) -> dict[str, list[str]]:
        keys: dict[str, list[str]] = {}
        for table_name, (_, primary_key, display_column) in BATCH_MODELS.items():
            ids = sorted(selection[table_name])
            if not ids:
                keys[table_name] = []
                continue
            result = await db.execute(select(display_column).where(primary_key.in_(ids)).order_by(display_column))
            keys[table_name] = sorted(str(row[0]) for row in result.all())
        return keys

    @classmethod
    async def delete_user_warehouse_links(
        cls,
        db: AsyncSession,
        warehouse_ids: Collection[int],
    ) -> None:
        if warehouse_ids:
            await db.execute(delete(BizUserWarehouse).where(BizUserWarehouse.warehouse_id.in_(sorted(warehouse_ids))))

    @classmethod
    async def delete_user_site_links(
        cls,
        db: AsyncSession,
        site_ids: Collection[int],
    ) -> None:
        if site_ids:
            await db.execute(delete(BizUserSite).where(BizUserSite.site_id.in_(sorted(site_ids))))

    @classmethod
    async def delete_table_rows(
        cls,
        db: AsyncSession,
        table_name: str,
        batch_id: int,
    ) -> int:
        if table_name not in BATCH_MODELS:
            raise ServiceException(message=f'不允许清理非安全表: {table_name}')
        model = BATCH_MODELS[table_name][0]
        result = await db.execute(delete(model).where(model.test_batch_id == batch_id))
        return int(result.rowcount or 0)