rarelens/api/app/routers/cases.py

import logging
import uuid
from dataclasses import asdict
from datetime import UTC, datetime

from fastapi import APIRouter, HTTPException, Query, status
from sqlalchemy import func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import selectinload

from app.config import settings
from app.db import SessionDep
from app.models import (
    Case,
    CasePhenotype,
    DecisionState,
    Job,
    JobStatus,
    Variant,
    VariantDecision,
)
from app.schemas import (
    CandidateOut,
    CandidatePage,
    CaseCreate,
    CaseOut,
    FunnelOut,
    JobOut,
    ProvenanceOut,
    ReportOut,
    ScoreOut,
)
from app.services import candidates as case_view
from app.services import events, triage
from app.services.scoring import score_job

logger = logging.getLogger(__name__)
router = APIRouter()

REPORT_TOP = 5


async def _case_or_404(session: SessionDep, case_id: uuid.UUID) -> Case:
    case = await case_view.get_case(session, case_id)
    if case is None:
        raise HTTPException(404, "case not found")
    return case


async def _shortlisted_counts(session: SessionDep, case_ids: list[uuid.UUID]) -> dict[uuid.UUID, int]:
    if not case_ids:
        return {}
    rows = await session.execute(
        select(Job.case_id, func.count(VariantDecision.id))
        .join(Variant, Variant.job_id == Job.id)
        .join(VariantDecision, VariantDecision.variant_id == Variant.id)
        .where(Job.case_id.in_(case_ids), VariantDecision.state == DecisionState.shortlisted)
        .group_by(Job.case_id)
    )
    return dict(rows.all())  # type: ignore[arg-type]


async def _latest_jobs(session: SessionDep, case_ids: list[uuid.UUID]) -> dict[uuid.UUID, Job]:
    if not case_ids:
        return {}
    jobs = await session.scalars(
        select(Job).where(Job.case_id.in_(case_ids)).order_by(Job.created_at.desc())
    )
    latest: dict[uuid.UUID, Job] = {}
    for job in jobs:  # ordered newest first, so the first one wins
        latest.setdefault(job.case_id, job)
    return latest


def _as_case_out(case: Case, job: Job | None, shortlisted: int) -> CaseOut:
    out = CaseOut.model_validate(case)
    out.latest_job = JobOut.model_validate(job) if job is not None else None
    out.shortlisted = shortlisted
    return out


@router.get("", response_model=list[CaseOut])
async def list_cases(session: SessionDep):
    cases = (
        await session.scalars(
            select(Case).order_by(Case.created_at.desc()).options(selectinload(Case.phenotypes))
        )
    ).all()
    ids = [c.id for c in cases]
    jobs, counts = await _latest_jobs(session, ids), await _shortlisted_counts(session, ids)
    return [_as_case_out(c, jobs.get(c.id), counts.get(c.id, 0)) for c in cases]


@router.post("", response_model=CaseOut, status_code=status.HTTP_201_CREATED)
async def create_case(payload: CaseCreate, session: SessionDep):
    case = Case(
        name=payload.name,
        vcf_uri=payload.vcf_uri,
        assembly=payload.assembly,
        phenotypes=[CasePhenotype(hpo_id=p.hpo_id, label=p.label) for p in payload.phenotypes],
    )
    session.add(case)
    try:
        await session.commit()
    except IntegrityError:  # cases.name is unique
        await session.rollback()
        raise HTTPException(409, f"a case named {payload.name!r} already exists") from None
    await session.refresh(case, attribute_names=["phenotypes"])
    return _as_case_out(case, None, 0)


@router.get("/{case_id}", response_model=CaseOut)
async def get_case(case_id: uuid.UUID, session: SessionDep):
    case = await _case_or_404(session, case_id)
    counts = await _shortlisted_counts(session, [case_id])
    return _as_case_out(case, await case_view.latest_job(session, case_id), counts.get(case_id, 0))


@router.get("/{case_id}/jobs", response_model=list[JobOut])
async def list_jobs(case_id: uuid.UUID, session: SessionDep):
    result = await session.scalars(
        select(Job).where(Job.case_id == case_id).order_by(Job.created_at.desc())
    )
    return result.all()


@router.post("/{case_id}/annotate", response_model=JobOut, status_code=status.HTTP_202_ACCEPTED)
async def annotate(case_id: uuid.UUID, session: SessionDep):
    case = await _case_or_404(session, case_id)
    # Commit `running` before launching: a local run that dies instantly is marked failed by its
    # watcher, and a later status write here would overwrite that.
    job = Job(case_id=case.id, status=JobStatus.running)
    session.add(job)
    await session.commit()
    try:
        job.workflow_ref = await events.launch(job.id, case.vcf_uri, case.assembly)
    except events.LaunchError as e:
        job.status = JobStatus.failed
        job.log = str(e)
        job.finished_at = datetime.now(UTC)
    await session.commit()
    await session.refresh(job)
    return job


@router.post("/{case_id}/score", response_model=ScoreOut)
async def score(case_id: uuid.UUID, session: SessionDep):
    await _case_or_404(session, case_id)
    job = await case_view.latest_job(session, case_id, status=JobStatus.succeeded)
    if job is None:
        raise HTTPException(409, "no finished annotation to score")
    try:
        scored, version = await score_job(job.id, session)
    except Exception as e:
        # No registry, no model behind the alias, a model that will not load: all of these are
        # the environment being unready, not a bug in the request. Name what is missing and keep
        # the exception in the log, where it is useful, rather than in the UI, where it is noise.
        logger.exception("scoring case %s failed", case_id)
        raise HTTPException(
            503,
            f"no model available: {settings.model_name}@{settings.model_alias} "
            f"via {settings.mlflow_tracking_uri}",
        ) from e
    return ScoreOut(case_id=case_id, scored=scored, model_version=version)


@router.get("/{case_id}/candidates", response_model=CandidatePage)
async def list_candidates(
    case_id: uuid.UUID,
    session: SessionDep,
    gene: str | None = None,
    impact: str | None = Query(None, pattern="^(HIGH|MODERATE|LOW|MODIFIER)$"),
    max_af: float | None = Query(None, ge=0, le=1),
    state: str | None = Query(None, pattern="^(shortlisted|dismissed|undecided)$"),
    limit: int = Query(50, ge=1, le=500),
    offset: int = Query(0, ge=0),
) -> CandidatePage:
    case = await _case_or_404(session, case_id)
    view = await case_view.build(session, case)

    items = view.candidates
    if gene:
        items = [c for c in items if (c.variant.gene or "").upper() == gene.upper()]
    if impact:
        items = [c for c in items if c.variant.impact == impact]
    if max_af is not None:
        items = [c for c in items if c.variant.gnomad_af is None or c.variant.gnomad_af <= max_af]
    if state == "undecided":
        items = [c for c in items if c.variant.decision is None]
    elif state is not None:
        items = [c for c in items if c.variant.decision and c.variant.decision.state.value == state]

    labels = view.labels
    return CandidatePage(
        # The funnel describes the whole case, not the filtered view.
        funnel=FunnelOut(**asdict(view.funnel)),
        weights=triage.WEIGHTS,
        items=[CandidateOut.from_candidate(c, labels) for c in items[offset : offset + limit]],
        total=len(items),
        limit=limit,
        offset=offset,
    )


@router.get("/{case_id}/report", response_model=ReportOut)
async def report(case_id: uuid.UUID, session: SessionDep) -> ReportOut:
    case = await _case_or_404(session, case_id)
    view = await case_view.build(session, case)
    labels = view.labels

    def decided(state: DecisionState) -> list[CandidateOut]:
        return [
            CandidateOut.from_candidate(c, labels)
            for c in view.candidates
            if c.variant.decision is not None and c.variant.decision.state == state
        ]

    shortlisted, dismissed = decided(DecisionState.shortlisted), decided(DecisionState.dismissed)
    top = [
        CandidateOut.from_candidate(c, labels)
        for c in view.candidates
        if c.variant.decision is None
    ][:REPORT_TOP]
    prediction = next((c.variant.prediction for c in view.candidates if c.variant.prediction), None)
    return ReportOut(
        case=_as_case_out(case, view.job, len(shortlisted)),
        funnel=FunnelOut(**asdict(view.funnel)),
        generated_at=datetime.now(UTC),
        provenance=ProvenanceOut(
            job_id=view.job.id if view.job else None,
            vep_version=view.job.vep_version if view.job else None,
            finished_at=view.job.finished_at if view.job else None,
            model_name=prediction.model_name if prediction else None,
            model_version=prediction.model_version if prediction else None,
        ),
        shortlisted=shortlisted,
        dismissed=dismissed,
        top=top,
    )