ctms/backend/app/crud/subject.py

import uuid
from datetime import date
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 _visit_code(visit) -> str:
    return str(getattr(visit, "visit_code", "") or "").strip()


def _is_cancelled_visit(visit) -> bool:
    return str(getattr(visit, "status", "") or "").strip().upper() == "CANCELLED"


def _standard_visit_codes(visit_schedule: list[dict] | None) -> set[str]:
    return {
        str(item.get("visit_code") or "").strip()
        for item in visit_schedule or []
        if str(item.get("visit_code") or "").strip()
        and str(item.get("visit_code") or "").strip() not in visit_crud.NON_TREATMENT_VISIT_CODES
    }


def derive_subject_status(subject, visits: Sequence[Visit], visit_schedule: list[dict] | None = None) -> str:
    if any(_visit_code(visit) == "提前终止" and getattr(visit, "actual_date", None) for visit in visits):
        return "DROPPED"

    standard_codes = _standard_visit_codes(visit_schedule)
    completion_visits = [
        visit
        for visit in visits
        if _visit_code(visit) not in visit_crud.NON_TREATMENT_VISIT_CODES
        and not _is_cancelled_visit(visit)
        and getattr(visit, "planned_date", None) is not None
    ]
    if standard_codes:
        standard_visits = [visit for visit in completion_visits if _visit_code(visit) in standard_codes]
        if standard_visits:
            latest_planned_date = max(getattr(visit, "planned_date") for visit in completion_visits)
            latest_visits = [
                visit for visit in completion_visits if getattr(visit, "planned_date", None) == latest_planned_date
            ]
            if latest_visits and all(getattr(visit, "actual_date", None) for visit in latest_visits):
                return "COMPLETED"
    elif completion_visits:
        latest_planned_date = max(getattr(visit, "planned_date") for visit in completion_visits)
        latest_visits = [visit for visit in completion_visits if getattr(visit, "planned_date", None) == latest_planned_date]
        if latest_visits and all(getattr(visit, "actual_date", None) for visit in latest_visits):
            return "COMPLETED"

    if getattr(subject, "enrollment_date", None):
        return "ENROLLED"
    return "SCREENING"


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("完成日期不能早于入组日期")


def _validate_actual_medication_count(value: int | None) -> None:
    if value is not None and value < 0:
        raise ValueError("实际用药次数不能小于0")


def _should_sync_visits(previous_baseline_date: date | None, next_baseline_date: date | None) -> bool:
    return next_baseline_date is not None and previous_baseline_date != next_baseline_date


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=subject_in.enrollment_date,
        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=subject_in.enrollment_date,
        baseline_date=subject_in.baseline_date,
        completion_date=None,
        drop_reason=None,
    )
    db.add(subject)
    await db.commit()
    await db.refresh(subject)

    await sync_visits_from_baseline(db, subject)
    await sync_subject_status(db, subject)
    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 sync_visits_from_baseline(db: AsyncSession, subject: Subject) -> None:
    # 基线/治疗日期是访视计划的唯一推算基准，不能用入组日期替代。
    result = await db.execute(select(Study).where(Study.id == subject.study_id))
    study = result.scalar_one_or_none()
    if not study:
        return

    await visit_crud.create_scheduled_visits(
        db,
        study_id=subject.study_id,
        subject=subject,
        base_date=subject.baseline_date,
        visit_schedule=study.visit_schedule,
    )


async def sync_subject_status(db: AsyncSession, subject: Subject) -> Subject:
    visits_result = await db.execute(select(Visit).where(Visit.subject_id == subject.id))
    visits = list(visits_result.scalars().all())
    study_result = await db.execute(select(Study.visit_schedule).where(Study.id == subject.study_id))
    visit_schedule = study_result.scalar_one_or_none() or []
    next_status = derive_subject_status(subject, visits, visit_schedule)
    if subject.status != next_status:
        await db.execute(
            sa_update(Subject)
            .where(Subject.id == subject.id)
            .values(status=next_status)
        )
        await db.commit()
        await db.refresh(subject)
    return subject


async def generate_default_visits(db: AsyncSession, subject: Subject) -> None:
    await sync_visits_from_baseline(db, subject)


async def update_subject(db: AsyncSession, subject: Subject, subject_in: SubjectUpdate) -> Subject:
    update_data = subject_in.model_dump(exclude_unset=True)
    update_data.pop("status", None)
    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,
    )
    _validate_actual_medication_count(update_data.get("actual_medication_count", subject.actual_medication_count))
    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 await sync_subject_status(db, subject)


def should_generate_visits_after_subject_update(
    *,
    previous_baseline_date: date | None,
    next_baseline_date: date | None,
) -> bool:
    return _should_sync_visits(previous_baseline_date, next_baseline_date)


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()