import uuid from dataclasses import dataclass 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.study import Study from app.models.subject import Subject from app.models.visit import Visit from app.schemas.visit import VisitCreate, VisitUpdate @dataclass(frozen=True) class EarlyTerminationVisitChanges: event_visit_code: str event_actual_date: date event_notes: str visits_to_cancel: list[Visit] NON_TREATMENT_VISIT_CODES = {"筛选访视", "基线访视", "提前终止"} def _visit_schedule_order_map(visit_schedule: list[dict] | None) -> dict[str, int]: order_map: dict[str, int] = {} for index, item in enumerate(visit_schedule or []): code = str(item.get("visit_code") or "").strip() if code and code not in order_map: order_map[code] = index return order_map def sort_visits_for_display(visits: Sequence[Visit], visit_schedule: list[dict] | None = None) -> list[Visit]: order_map = _visit_schedule_order_map(visit_schedule) fallback_start = len(order_map) early_termination = next((visit for visit in visits if (visit.visit_code or "").strip() == "提前终止"), None) if early_termination and early_termination.actual_date: last_actual_index = max( ( order_map.get((visit.visit_code or "").strip(), fallback_start) for visit in visits if visit is not early_termination and visit.actual_date is not None ), default=-1, ) sorted_visits: list[Visit] = [] inserted = False for visit in sorted( [visit for visit in visits if visit is not early_termination], key=lambda visit: ( order_map.get((visit.visit_code or "").strip(), fallback_start), visit.planned_date is None, visit.planned_date or date.max, visit.visit_code or "", ), ): sorted_visits.append(visit) if not inserted and order_map.get((visit.visit_code or "").strip(), fallback_start) == last_actual_index: sorted_visits.append(early_termination) inserted = True if not inserted: sorted_visits.insert(0, early_termination) return sorted_visits return sorted( visits, key=lambda visit: ( order_map.get((visit.visit_code or "").strip(), fallback_start), visit.planned_date is None, visit.planned_date or date.max, visit.visit_code or "", ), ) def build_visit_schedule_dates(visit_schedule: list[dict] | None, base_date: date | None) -> list[dict]: if not visit_schedule: return [] rows: list[dict] = [] for item in visit_schedule: code = str(item.get("visit_code") or "").strip() if not code: continue baseline_offset_days = int(item.get("baseline_offset_days", 0)) window_before_days = int(item.get("window_before_days", 0)) window_after_days = int(item.get("window_after_days", 0)) planned_date = base_date + timedelta(days=baseline_offset_days) if base_date else None rows.append( { "visit_code": code, "baseline_offset_days": baseline_offset_days, "planned_date": planned_date, "window_start": planned_date - timedelta(days=window_before_days) if planned_date else None, "window_end": planned_date + timedelta(days=window_after_days) if planned_date else None, } ) return rows def get_visit_window_start_date(visit: Visit) -> date | None: return visit.window_start or visit.planned_date def get_last_planned_visit_window_start_date(visits: Sequence[Visit]) -> date | None: window_start_dates = [ get_visit_window_start_date(visit) for visit in visits if get_visit_window_start_date(visit) is not None and (visit.visit_code or "").strip() not in NON_TREATMENT_VISIT_CODES ] return max(window_start_dates) if window_start_dates else None def validate_early_termination_date(termination_date: date, visits: Sequence[Visit]) -> None: last_window_start_date = get_last_planned_visit_window_start_date(visits) if last_window_start_date is not None and termination_date >= last_window_start_date: raise ValueError(f"提前终止日期必须早于方案最后一个计划访视窗口开始日({last_window_start_date.isoformat()})") def build_early_termination_visit_changes( visits: Sequence[Visit], *, termination_date: date, reason: str, notes: str | None = None, ) -> EarlyTerminationVisitChanges: event_notes = reason.strip() extra_notes = (notes or "").strip() if extra_notes: event_notes = f"{event_notes}\n{extra_notes}" visits_to_cancel = [ visit for visit in visits if visit.actual_date is None and visit.status not in ["DONE", "CANCELLED", "LOST"] and (visit.visit_code or "").strip() not in NON_TREATMENT_VISIT_CODES ] return EarlyTerminationVisitChanges( event_visit_code="提前终止", event_actual_date=termination_date, event_notes=event_notes, visits_to_cancel=visits_to_cancel, ) 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, actual_date: date | None = None, status: str = "PLANNED", ) -> Visit: visit = Visit( study_id=study_id, subject_id=subject.id, visit_code=visit_code, planned_date=planned_date, actual_date=actual_date, status=status, 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)) visits = result.scalars().all() if not visits: return [] study_result = await db.execute(select(Study.visit_schedule).where(Study.id == visits[0].study_id)) visit_schedule = study_result.scalar_one_or_none() or [] return sort_visits_for_display(visits, visit_schedule) 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_scheduled_visits( db: AsyncSession, *, study_id: uuid.UUID, subject: Subject, base_date: date | None, visit_schedule: list[dict] | None, ) -> Sequence[Visit]: if not visit_schedule: return [] result = await db.execute(select(Visit).where(Visit.subject_id == subject.id)) existing_by_code = {visit.visit_code: visit for visit in result.scalars().all()} created: list[Visit] = [] for item in build_visit_schedule_dates(visit_schedule, base_date): code = item["visit_code"] is_baseline_visit = item["baseline_offset_days"] == 0 and base_date is not None existing = existing_by_code.get(code) if existing: if is_baseline_visit: await db.execute( sa_update(Visit) .where(Visit.id == existing.id) .values( planned_date=item["planned_date"], window_start=item["window_start"], window_end=item["window_end"], actual_date=base_date, status="DONE", ) ) await db.commit() elif existing.status == "PLANNED" and existing.actual_date is None: await db.execute( sa_update(Visit) .where(Visit.id == existing.id) .values( planned_date=item["planned_date"], window_start=item["window_start"], window_end=item["window_end"], ) ) await db.commit() continue created.append( await create_visit( db, study_id=study_id, visit_in=None, subject=subject, visit_code=code, planned_date=item["planned_date"], window_start=item["window_start"], window_end=item["window_end"], actual_date=base_date if is_baseline_visit else None, status="DONE" if is_baseline_visit else "PLANNED", ) ) return created async def create_early_termination_visit( db: AsyncSession, *, study_id: uuid.UUID, subject: Subject, termination_date: date, reason: str, notes: str | None = None, ) -> Visit: result = await db.execute(select(Visit).where(Visit.subject_id == subject.id)) existing_visits = list(result.scalars().all()) validate_early_termination_date(termination_date, existing_visits) changes = build_early_termination_visit_changes( existing_visits, termination_date=termination_date, reason=reason, notes=notes, ) existing_event = next( (visit for visit in existing_visits if visit.visit_code == changes.event_visit_code), None, ) if existing_event: await db.execute( sa_update(Visit) .where(Visit.id == existing_event.id) .values( actual_date=changes.event_actual_date, status="DONE", notes=changes.event_notes, ) ) event_visit = existing_event else: event_visit = Visit( study_id=study_id, subject_id=subject.id, visit_code=changes.event_visit_code, planned_date=None, actual_date=changes.event_actual_date, status="DONE", window_start=None, window_end=None, notes=changes.event_notes, ) db.add(event_visit) for visit in changes.visits_to_cancel: await db.execute( sa_update(Visit) .where(Visit.id == visit.id) .values(status="CANCELLED", notes="提前终止后不再适用") ) await db.commit() await db.refresh(event_visit) return event_visit