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