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

from exceptions.exception import ServiceWarning
from module_business.domain.test_batch_graph import BATCH_MODELS, ROOT_TABLES, BatchSelection
from module_business.entity.vo.test_batch_vo import AdoptionRootModel


class TestBatchAdoptionDao:
    """旧数据闭包锁定、展示与批量纳入。"""

    @classmethod
    async def lock_roots(
        cls,
        db: AsyncSession,
        roots: list[AdoptionRootModel],
    ) -> None:
        requested: dict[str, set[int]] = {}
        for root in roots:
            requested.setdefault(ROOT_TABLES[root.resource_type], set()).update(root.resource_ids)
        for table_name in sorted(requested):
            primary_key = BATCH_MODELS[table_name][1]
            ids = sorted(requested[table_name])
            result = await db.execute(
                select(primary_key).where(primary_key.in_(ids)).order_by(primary_key).with_for_update()
            )
            if {row[0] for row in result.all()} != set(ids):
                raise ServiceWarning('', '指定的旧数据不存在')

    @classmethod
    async def lock_selection(cls, db: AsyncSession, selection: BatchSelection) -> None:
        for table_name in sorted(selection):
            ids = sorted(selection[table_name])
            if not ids:
                continue
            primary_key = BATCH_MODELS[table_name][1]
            result = await db.execute(
                select(primary_key).where(primary_key.in_(ids)).order_by(primary_key).with_for_update()
            )
            if len(result.all()) != len(ids):
                raise ServiceWarning('', '预览范围已变化')

    @classmethod
    async def find_tagged_record(
        cls,
        db: AsyncSession,
        selection: BatchSelection,
    ) -> tuple[str, int] | None:
        for table_name in sorted(selection):
            ids = sorted(selection[table_name])
            if not ids:
                continue
            model, primary_key, _ = BATCH_MODELS[table_name]
            result = await db.execute(
                select(primary_key)
                .where(
                    primary_key.in_(ids),
                    model.test_batch_id.is_not(None),
                )
                .order_by(primary_key)
                .limit(1)
            )
            record_id = result.scalar_one_or_none()
            if record_id is not None:
                return table_name, record_id
        return None

    @classmethod
    async def list_display_keys(
        cls,
        db: AsyncSession,
        selection: BatchSelection,
    ) -> dict[str, list[str]]:
        keys: dict[str, list[str]] = {}
        for table_name in BATCH_MODELS:
            ids = sorted(selection[table_name])
            if not ids:
                keys[table_name] = []
                continue
            _, primary_key, display_column = BATCH_MODELS[table_name]
            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 set_test_batch_id(
        cls,
        db: AsyncSession,
        table_name: str,
        ids: list[int],
        batch_id: int,
    ) -> None:
        if not ids:
            return
        model, primary_key, _ = BATCH_MODELS[table_name]
        result = await db.execute(
            update(model).where(primary_key.in_(ids), model.test_batch_id.is_(None)).values(test_batch_id=batch_id)
        )
        if result.rowcount != len(ids):
            raise ServiceWarning('', '存在记录已属于其他测试批次')