ctms/backend/app/crud/visit.py

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

from sqlalchemy import and_, func, or_, select, update as sa_update
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.subject import Subject
from app.models.visit import Visit
from app.schemas.visit import VisitCreate, VisitUpdate


async def create_visit(
    db: AsyncSession,
    *,
    study_id: uuid.UUID,
    visit_in: VisitCreate | None,
    subject: Subject,
    visit_code: str,
    planned_date,
    window_start: date | None = None,
    window_end: date | None = None,
) -> Visit:
    visit = Visit(
        study_id=study_id,
        subject_id=subject.id,
        visit_code=visit_code,
        planned_date=planned_date,
        actual_date=None,
        status="PLANNED",
        window_start=window_start if window_start is not None else (visit_in.window_start if visit_in else None),
        window_end=window_end if window_end is not None else (visit_in.window_end if visit_in else None),
        notes=None,
    )
    db.add(visit)
    await db.commit()
    await db.refresh(visit)
    return visit


async def list_visits(db: AsyncSession, subject_id: uuid.UUID) -> Sequence[Visit]:
    result = await db.execute(select(Visit).where(Visit.subject_id == subject_id).order_by(Visit.planned_date))
    return result.scalars().all()


async def mark_overdue_as_lost(db: AsyncSession, subject_id: uuid.UUID) -> None:
    today = func.current_date()
    deadline = func.coalesce(Visit.window_end, Visit.planned_date)
    await db.execute(
        sa_update(Visit)
        .where(
            Visit.subject_id == subject_id,
            Visit.actual_date.is_(None),
            deadline.is_not(None),
            deadline < today,
            Visit.status.not_in(["DONE", "CANCELLED", "LOST"]),
        )
        .values(status="LOST")
    )
    await db.commit()


async def mark_overdue_as_lost_global(db: AsyncSession) -> None:
    today = func.current_date()
    deadline = func.coalesce(Visit.window_end, Visit.planned_date)
    await db.execute(
        sa_update(Visit)
        .where(
            Visit.actual_date.is_(None),
            deadline.is_not(None),
            deadline < today,
            Visit.status.not_in(["DONE", "CANCELLED", "LOST"]),
        )
        .values(status="LOST")
    )
    await db.commit()


async def get_next_visit_code(db: AsyncSession, subject_id: uuid.UUID) -> str:
    result = await db.execute(select(Visit.visit_code).where(Visit.subject_id == subject_id))
    used: set[int] = set()
    for (code,) in result.all():
        if not code or not code.startswith("V"):
            continue
        num = code[1:]
        if num.isdigit():
            used.add(int(num))
    next_num = 1
    while next_num in used:
        next_num += 1
    return f"V{next_num}"


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


async def update_visit(db: AsyncSession, visit: Visit, visit_in: VisitUpdate) -> Visit:
    update_data = visit_in.model_dump(exclude_unset=True)
    if "actual_date" in update_data and update_data["actual_date"] is not None and "status" not in update_data:
        actual_date = update_data["actual_date"]
        if (visit.window_start and actual_date < visit.window_start) or (visit.window_end and actual_date > visit.window_end):
            update_data["status"] = "OVERDUE"
        else:
            update_data["status"] = "DONE"
    if update_data:
        await db.execute(
            sa_update(Visit)
            .where(Visit.id == visit.id)
            .values(**update_data)
        )
        await db.commit()
    await db.refresh(visit)
    return visit


async def delete_visit(db: AsyncSession, visit: Visit) -> None:
    await db.delete(visit)
    await db.commit()


async def list_lost_visits(
    db: AsyncSession,
    study_id: uuid.UUID,
    *,
    site_ids: set[uuid.UUID] | None = None,
    limit: int = 50,
) -> Sequence[tuple[Visit, str, uuid.UUID | None]]:
    if site_ids is not None and not site_ids:
        return []
    today = func.current_date()
    deadline = func.coalesce(Visit.window_end, Visit.planned_date)
    lost_condition = and_(
        Visit.actual_date.is_(None),
        deadline.is_not(None),
        deadline < today,
        Visit.status.not_in(["DONE", "CANCELLED"]),
    )
    stmt = (
        select(Visit, Subject.subject_no, Subject.site_id)
        .join(Subject, Subject.id == Visit.subject_id)
        .where(Visit.study_id == study_id, lost_condition)
        .order_by(Visit.updated_at.desc())
        .limit(limit)
    )
    if site_ids is not None:
        stmt = stmt.where(Subject.site_id.in_(site_ids))
    result = await db.execute(stmt)
    return result.all()


async def create_followup_visits(
    db: AsyncSession,
    *,
    study_id: uuid.UUID,
    subject: Subject,
    base_date: date,
    visit_total: int | None,
    visit_interval_days: int | None,
    window_start_offset: int | None,
    window_end_offset: int | None,
) -> Sequence[Visit]:
    if not visit_total or not visit_interval_days or visit_total < 2:
        return []

    result = await db.execute(select(Visit.visit_code).where(Visit.subject_id == subject.id))
    existing_codes = {row[0] for row in result.all()}
    created: list[Visit] = []
    for index in range(2, visit_total + 1):
        code = f"V{index}"
        if code in existing_codes:
            continue
        planned_date = base_date + timedelta(days=visit_interval_days * (index - 1))
        window_start = (
            planned_date + timedelta(days=window_start_offset) if window_start_offset is not None else None
        )
        window_end = planned_date + timedelta(days=window_end_offset) if window_end_offset is not None else None
        created.append(
            await create_visit(
                db,
                study_id=study_id,
                visit_in=None,
                subject=subject,
                visit_code=code,
                planned_date=planned_date,
                window_start=window_start,
                window_end=window_end,
            )
        )
    return created