ctms/backend/app/crud/visit.py

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 derive_visit_status_from_actual_date(
    *,
    actual_date: date | None,
    window_start: date | None,
    window_end: date | None,
    default_status: str = "PLANNED",
) -> str:
    if actual_date is None:
        return default_status
    if (window_start and actual_date < window_start) or (window_end and actual_date > window_end):
        return "OVERDUE"
    return "DONE"


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:
    resolved_window_start = window_start if window_start is not None else (visit_in.window_start if visit_in else None)
    resolved_window_end = window_end if window_end is not None else (visit_in.window_end if visit_in else None)
    resolved_actual_date = actual_date if actual_date is not None else (visit_in.actual_date if visit_in else None)
    visit = Visit(
        study_id=study_id,
        subject_id=subject.id,
        visit_code=visit_code,
        planned_date=planned_date,
        actual_date=resolved_actual_date,
        status=derive_visit_status_from_actual_date(
            actual_date=resolved_actual_date,
            window_start=resolved_window_start,
            window_end=resolved_window_end,
            default_status=status,
        ),
        window_start=resolved_window_start,
        window_end=resolved_window_end,
        notes=visit_in.notes if visit_in else 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:
        update_data["status"] = derive_visit_status_from_actual_date(
            actual_date=update_data["actual_date"],
            window_start=visit.window_start,
            window_end=visit.window_end,
        )
    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