ctms/backend/app/crud/subject.py

import uuid
from datetime import date, timedelta
from typing import Sequence

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

from app.crud import visit as visit_crud
from app.models.ae import AdverseEvent
from app.models.finance import FinanceItem
from app.models.study import Study
from app.models.subject_history import SubjectHistory
from app.models.visit import Visit
from app.models.site import Site
from app.models.subject import Subject
from app.schemas.subject import SubjectCreate, SubjectUpdate


def _validate_subject_date_chain(
    *,
    screening_date: date | None,
    consent_date: date | None,
    enrollment_date: date | None,
    completion_date: date | None,
) -> None:
    if screening_date and consent_date and consent_date < screening_date:
        raise ValueError("知情同意日期不能早于筛选日期")
    if screening_date and enrollment_date and enrollment_date < screening_date:
        raise ValueError("入组日期不能早于筛选日期")
    if consent_date and enrollment_date and enrollment_date < consent_date:
        raise ValueError("入组日期不能早于知情同意日期")
    if screening_date and completion_date and completion_date < screening_date:
        raise ValueError("完成日期不能早于筛选日期")
    if consent_date and completion_date and completion_date < consent_date:
        raise ValueError("完成日期不能早于知情同意日期")
    if enrollment_date and completion_date and completion_date < enrollment_date:
        raise ValueError("完成日期不能早于入组日期")


async def _validate_site(db: AsyncSession, study_id: uuid.UUID, site_id: uuid.UUID) -> None:
    result = await db.execute(select(Site).where(Site.id == site_id))
    site = result.scalar_one_or_none()
    if not site or site.study_id != study_id:
        raise ValueError("分中心不属于当前项目")
    if not site.is_active:
        raise ValueError("中心已停用")


async def create_subject(db: AsyncSession, study_id: uuid.UUID, subject_in: SubjectCreate) -> Subject:
    await _validate_site(db, study_id, subject_in.site_id)
    _validate_subject_date_chain(
        screening_date=subject_in.screening_date,
        consent_date=subject_in.consent_date,
        enrollment_date=None,
        completion_date=None,
    )
    subject = Subject(
        study_id=study_id,
        site_id=subject_in.site_id,
        subject_no=subject_in.subject_no,
        status="SCREENING",
        screening_date=subject_in.screening_date,
        consent_date=subject_in.consent_date,
        enrollment_date=None,
        completion_date=None,
        drop_reason=None,
    )
    db.add(subject)
    await db.commit()
    await db.refresh(subject)

    # initial visit: Screening (V0)
    await visit_crud.create_visit(
        db,
        study_id=study_id,
        visit_in=None,
        subject=subject,
        visit_code="V0",
        planned_date=subject.screening_date,
    )
    return subject


async def get_subject(db: AsyncSession, subject_id: uuid.UUID) -> Subject | None:
    result = await db.execute(select(Subject).where(Subject.id == subject_id))
    return result.scalar_one_or_none()


async def list_subjects(
    db: AsyncSession,
    study_id: uuid.UUID,
    site_id: uuid.UUID | None = None,
    site_ids: set[uuid.UUID] | None = None,
) -> Sequence[Subject]:
    stmt = select(Subject).where(Subject.study_id == study_id)
    if site_ids is not None:
        if not site_ids:
            return []
        stmt = stmt.where(Subject.site_id.in_(site_ids))
    if site_id:
        stmt = stmt.where(Subject.site_id == site_id)
    result = await db.execute(stmt)
    return result.scalars().all()


async def generate_default_visits(db: AsyncSession, subject: Subject) -> None:
    # Baseline + Follow-up visits based on enrollment_date
    if not subject.enrollment_date:
        return
    result = await db.execute(select(Study).where(Study.id == subject.study_id))
    study = result.scalar_one_or_none()
    if not study:
        return

    visit_total = study.visit_total or 3
    visit_interval_days = study.visit_interval_days or 30
    window_start_offset = study.visit_window_start_offset
    window_end_offset = study.visit_window_end_offset
    baseline_date = subject.enrollment_date

    result = await db.execute(select(Visit.visit_code).where(Visit.subject_id == subject.id))
    existing_codes = {row[0] for row in result.all()}
    if "V1" not in existing_codes:
        window_start = baseline_date + timedelta(days=window_start_offset) if window_start_offset is not None else None
        window_end = baseline_date + timedelta(days=window_end_offset) if window_end_offset is not None else None
        await visit_crud.create_visit(
            db,
            study_id=subject.study_id,
            visit_in=None,
            subject=subject,
            visit_code="V1",
            planned_date=baseline_date,
            window_start=window_start,
            window_end=window_end,
        )

    await visit_crud.create_followup_visits(
        db,
        study_id=subject.study_id,
        subject=subject,
        base_date=baseline_date,
        visit_total=visit_total,
        visit_interval_days=visit_interval_days,
        window_start_offset=window_start_offset,
        window_end_offset=window_end_offset,
    )


async def update_subject(db: AsyncSession, subject: Subject, subject_in: SubjectUpdate) -> Subject:
    update_data = subject_in.model_dump(exclude_unset=True)
    next_screening_date = subject.screening_date
    next_consent_date = update_data.get("consent_date", subject.consent_date)
    next_enrollment_date = update_data.get("enrollment_date", subject.enrollment_date)
    next_completion_date = update_data.get("completion_date", subject.completion_date)
    _validate_subject_date_chain(
        screening_date=next_screening_date,
        consent_date=next_consent_date,
        enrollment_date=next_enrollment_date,
        completion_date=next_completion_date,
    )
    if update_data:
        await db.execute(
            sa_update(Subject)
            .where(Subject.id == subject.id)
            .values(**update_data)
        )
        await db.commit()
    await db.refresh(subject)
    return subject


async def delete_subject(db: AsyncSession, subject: Subject) -> None:
    subject_id = subject.id
    await db.execute(sa_delete(Visit).where(Visit.subject_id == subject_id))
    await db.execute(sa_delete(SubjectHistory).where(SubjectHistory.subject_id == subject_id))
    await db.execute(sa_delete(AdverseEvent).where(AdverseEvent.subject_id == subject_id))
    await db.execute(sa_delete(FinanceItem).where(FinanceItem.subject_id == subject_id))
    await db.delete(subject)
    await db.commit()