ctms/backend/app/crud/study_setup_config.py

import uuid
from datetime import datetime

from sqlalchemy import delete as sa_delete, select
from sqlalchemy import func
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.study_setup_config import StudySetupConfig
from app.models.study_setup_config_version import StudySetupConfigVersion
from app.schemas.study_setup_config import StudySetupConfigData


def _to_storage(data: StudySetupConfigData | dict) -> dict:
    if isinstance(data, StudySetupConfigData):
        return data.model_dump(mode="json")
    return data


async def get_by_study(db: AsyncSession, study_id: uuid.UUID) -> StudySetupConfig | None:
    result = await db.execute(select(StudySetupConfig).where(StudySetupConfig.study_id == study_id))
    return result.scalar_one_or_none()


async def upsert(
    db: AsyncSession,
    study_id: uuid.UUID,
    *,
    expected_version: int | None,
    data: StudySetupConfigData | dict,
    saved_by: uuid.UUID | None,
) -> tuple[StudySetupConfig | None, bool]:
    existing = await get_by_study(db, study_id)
    payload = _to_storage(data)
    if existing:
        if expected_version is not None and expected_version != existing.version:
            return None, True
        existing.version = existing.version + 1
        existing.config = payload
        existing.saved_by = saved_by
        existing.publish_status = "PUBLISHED" if existing.published_config == payload else "DRAFT"
        existing.updated_at = datetime.utcnow()
        await db.commit()
        await db.refresh(existing)
        return existing, False

    record = StudySetupConfig(
        study_id=study_id,
        version=1,
        config=payload,
        publish_status="DRAFT",
        saved_by=saved_by,
    )
    db.add(record)
    await db.commit()
    await db.refresh(record)
    return record, False


async def publish(
    db: AsyncSession,
    study_id: uuid.UUID,
    *,
    expected_version: int | None,
    published_by: uuid.UUID | None,
    auto_commit: bool = True,
) -> tuple[StudySetupConfig | None, bool]:
    existing = await get_by_study(db, study_id)
    if not existing:
        return None, False
    if expected_version is not None and expected_version != existing.version:
        return None, True
    next_version = existing.version + 1
    next_display_version_stmt = select(func.max(StudySetupConfigVersion.display_version)).where(
        StudySetupConfigVersion.study_id == study_id
    )
    display_version_result = await db.execute(next_display_version_stmt)
    current_max_display_version = display_version_result.scalar_one_or_none() or 0
    next_display_version = current_max_display_version + 1
    existing.version = next_version
    existing.publish_status = "PUBLISHED"
    existing.published_config = existing.config
    existing.published_by = published_by
    existing.published_at = datetime.utcnow()
    existing.saved_by = published_by
    existing.updated_at = datetime.utcnow()
    snapshot = StudySetupConfigVersion(
        study_setup_config_id=existing.id,
        study_id=study_id,
        version=next_version,
        display_version=next_display_version,
        source_version=next_version - 1,
        config=existing.config,
        published_by=published_by,
        published_at=existing.published_at,
    )
    db.add(snapshot)
    if auto_commit:
        await db.commit()
        await db.refresh(existing)
    else:
        await db.flush()
    return existing, False


async def list_versions(db: AsyncSession, study_id: uuid.UUID) -> list[StudySetupConfigVersion]:
    result = await db.execute(
        select(StudySetupConfigVersion)
        .where(StudySetupConfigVersion.study_id == study_id)
        .order_by(StudySetupConfigVersion.version.desc(), StudySetupConfigVersion.published_at.desc())
    )
    return list(result.scalars().all())


async def rollback_to_version(
    db: AsyncSession,
    study_id: uuid.UUID,
    *,
    expected_version: int | None,
    target_version: int,
    saved_by: uuid.UUID | None,
) -> tuple[StudySetupConfig | None, bool, bool]:
    existing = await get_by_study(db, study_id)
    if not existing:
        return None, False, False
    if expected_version is not None and expected_version != existing.version:
        return None, True, False

    result = await db.execute(
        select(StudySetupConfigVersion).where(
            StudySetupConfigVersion.study_id == study_id,
            StudySetupConfigVersion.version == target_version,
        )
    )
    target = result.scalar_one_or_none()
    if not target:
        return None, False, True

    existing.version = existing.version + 1
    existing.config = target.config
    existing.saved_by = saved_by
    existing.publish_status = "DRAFT"
    existing.updated_at = datetime.utcnow()
    await db.commit()
    await db.refresh(existing)
    return existing, False, False


async def delete_version(
    db: AsyncSession,
    study_id: uuid.UUID,
    *,
    target_version: int,
) -> bool:
    result = await db.execute(
        select(StudySetupConfigVersion).where(
            StudySetupConfigVersion.study_id == study_id,
            StudySetupConfigVersion.version == target_version,
        )
    )
    target = result.scalar_one_or_none()
    if not target:
        return False
    await db.delete(target)
    await db.commit()
    return True


async def delete_by_study(db: AsyncSession, study_id: uuid.UUID) -> None:
    await db.execute(sa_delete(StudySetupConfigVersion).where(StudySetupConfigVersion.study_id == study_id))
    await db.execute(sa_delete(StudySetupConfig).where(StudySetupConfig.study_id == study_id))
    await db.commit()