ctms/backend/app/crud/study_setup_config.py

import uuid
from copy import deepcopy
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 deepcopy(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,
    force_draft: bool = False,
) -> 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
        if force_draft:
            existing.publish_status = "DRAFT"
        elif existing.publish_status == "DRAFT":
            # Keep draft state sticky until explicit publish API is called.
            existing.publish_status = "DRAFT"
        else:
            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
    publish_payload = _to_storage(existing.config or {})
    next_version = existing.version + 1
    latest_snapshot_stmt = (
        select(StudySetupConfigVersion)
        .where(StudySetupConfigVersion.study_id == study_id)
        .order_by(StudySetupConfigVersion.version.desc(), StudySetupConfigVersion.published_at.desc())
        .limit(1)
    )
    latest_snapshot_result = await db.execute(latest_snapshot_stmt)
    latest_snapshot = latest_snapshot_result.scalar_one_or_none()
    should_create_snapshot = not latest_snapshot or (latest_snapshot.config or {}) != publish_payload

    existing.version = next_version
    existing.publish_status = "PUBLISHED"
    existing.published_config = _to_storage(publish_payload)
    existing.published_by = published_by
    existing.published_at = datetime.utcnow()
    existing.saved_by = published_by
    existing.updated_at = datetime.utcnow()

    if should_create_snapshot:
        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
        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=_to_storage(publish_payload),
            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 = _to_storage(target.config or {})
    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()