rarelens/api/app/schemas.py

import re
import uuid
from datetime import datetime
from pathlib import PurePosixPath
from typing import TYPE_CHECKING, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator

from app.config import settings
from app.models import DecisionState, JobStatus

if TYPE_CHECKING:
    from app.services.triage import Candidate


class ORMModel(BaseModel):
    model_config = ConfigDict(from_attributes=True)


Assembly = Literal["GRCh38", "GRCh37"]
VCF_SUFFIXES = (".vcf", ".vcf.gz", ".vcf.bgz", ".bcf")
GCS_URI = re.compile(r"gs://[a-z0-9][a-z0-9._-]{1,220}[a-z0-9]/\S+")


class PhenotypeTerm(ORMModel):
    """An HPO term: the id is what ranking matches on, the label is for people."""

    hpo_id: str = Field(pattern=r"^HP:\d{7}$")
    label: str = Field(min_length=1, max_length=200)


class CaseCreate(BaseModel):
    name: str = Field(min_length=1, max_length=120)
    vcf_uri: str
    assembly: Assembly = "GRCh38"
    phenotypes: list[PhenotypeTerm] = Field(default_factory=list, max_length=100)

    @field_validator("vcf_uri")
    @classmethod
    def _gcs_object_or_file_under_data_root(cls, v: str) -> str:
        # The URI becomes a Nextflow argument and a path the pipeline reads: accept a GCS object or
        # a file under the local data root, never something that parses as an option.
        if not v or any(ord(c) < 32 for c in v):
            raise ValueError("vcf_uri must be a non-empty single line")
        if not v.lower().endswith(VCF_SUFFIXES):
            raise ValueError(f"vcf_uri must end in one of {', '.join(VCF_SUFFIXES)}")
        if v.startswith("gs://"):
            if not GCS_URI.fullmatch(v):
                raise ValueError("vcf_uri is not a valid gs://bucket/object URI")
            return v
        path, root = PurePosixPath(v), PurePosixPath(settings.local_data_root)
        if not path.is_absolute() or ".." in path.parts or not path.is_relative_to(root):
            raise ValueError(f"local VCFs must be absolute paths under {root}")
        return v


class JobOut(ORMModel):
    id: uuid.UUID
    case_id: uuid.UUID
    status: JobStatus
    workflow_ref: str | None
    vep_version: str | None
    log: str | None
    created_at: datetime
    finished_at: datetime | None


class CaseOut(ORMModel):
    id: uuid.UUID
    name: str
    vcf_uri: str
    assembly: str
    created_at: datetime
    phenotypes: list[PhenotypeTerm] = Field(default_factory=list)
    latest_job: JobOut | None = None
    shortlisted: int = 0


class PredictionOut(ORMModel):
    model_name: str
    model_version: str
    score: float


class VariantOut(ORMModel):
    id: int
    chrom: str
    pos: int
    ref: str
    alt: str
    gene: str | None
    consequence: str | None
    impact: str | None
    hgvsc: str | None
    hgvsp: str | None
    gnomad_af: float | None
    clinvar_sig: str | None
    prediction: PredictionOut | None = None


class DecisionIn(BaseModel):
    state: DecisionState
    reason: str | None = Field(None, max_length=120)
    note: str | None = Field(None, max_length=2000)


class DecisionOut(ORMModel):
    state: DecisionState
    reason: str | None
    note: str | None
    decided_at: datetime


class CandidateOut(BaseModel):
    variant: VariantOut
    score: float
    components: dict[str, float]
    matched_terms: list[PhenotypeTerm]
    scored: bool
    decision: DecisionOut | None = None

    @classmethod
    def from_candidate(cls, candidate: "Candidate", labels: dict[str, str]) -> "CandidateOut":
        variant = candidate.variant
        return cls(
            variant=VariantOut.model_validate(variant),
            score=round(candidate.score, 4),
            components={name: round(v, 4) for name, v in candidate.components.items()},
            matched_terms=[
                PhenotypeTerm(hpo_id=term, label=labels.get(term, term))
                for term in candidate.matched_terms
            ],
            scored=candidate.scored,
            decision=DecisionOut.model_validate(variant.decision) if variant.decision else None,
        )


class VariantDetailOut(CandidateOut):
    annotations: dict


class FunnelOut(BaseModel):
    total: int
    rare: int
    candidates: int
    phenotype_matched: int


class CandidatePage(BaseModel):
    funnel: FunnelOut
    weights: dict[str, float]
    items: list[CandidateOut]
    total: int
    limit: int
    offset: int


class ProvenanceOut(BaseModel):
    job_id: uuid.UUID | None = None
    vep_version: str | None = None
    finished_at: datetime | None = None
    model_name: str | None = None
    model_version: str | None = None


class ReportOut(BaseModel):
    case: CaseOut
    funnel: FunnelOut
    generated_at: datetime
    provenance: ProvenanceOut
    shortlisted: list[CandidateOut]
    dismissed: list[CandidateOut]
    # What a reviewer would look at next; a report with no decisions yet still says something.
    top: list[CandidateOut]


class ScoreOut(BaseModel):
    case_id: uuid.UUID
    scored: int
    model_version: str