"""Validation and loading helpers for image fixtures in test definitions."""

from __future__ import annotations

import asyncio

from db.models import TestAsset
from services.test_assets import read_test_image
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from test_schemas import TestCaseDefinition


def image_asset_ids(definition: TestCaseDefinition) -> set[str]:
    return {
        turn.image.asset_id
        for turn in definition.turns
        if turn.image is not None
    }


async def canonicalize_test_assets(
    session: AsyncSession,
    definition: TestCaseDefinition,
) -> None:
    """Validate references and replace client metadata with authoritative values."""

    asset_ids = image_asset_ids(definition)
    if not asset_ids:
        return
    rows = (
        await session.execute(select(TestAsset).where(TestAsset.id.in_(asset_ids)))
    ).scalars().all()
    by_id = {row.id: row for row in rows}
    missing = sorted(asset_ids - by_id.keys())
    if missing:
        raise ValueError("测试图片不存在: " + "、".join(missing))

    for turn in definition.turns:
        if turn.image is None:
            continue
        row = by_id[turn.image.asset_id]
        turn.image.filename = row.original_name
        turn.image.width = row.width
        turn.image.height = row.height
        turn.image.size_bytes = row.size_bytes


async def load_test_asset_bytes(
    session: AsyncSession,
    definition: TestCaseDefinition,
) -> dict[str, bytes]:
    """Load immutable image fixtures once before a transportless test attempt."""

    asset_ids = image_asset_ids(definition)
    if not asset_ids:
        return {}
    rows = (
        await session.execute(select(TestAsset).where(TestAsset.id.in_(asset_ids)))
    ).scalars().all()
    by_id = {row.id: row for row in rows}
    missing = sorted(asset_ids - by_id.keys())
    if missing:
        raise ValueError("测试图片不存在: " + "、".join(missing))

    pairs = await asyncio.gather(
        *(
            asyncio.to_thread(read_test_image, row.storage_key)
            for row in rows
        )
    )
    return {row.id: data for row, data in zip(rows, pairs, strict=True)}