src/schemas/report.py

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator

from src.schemas.detection import LocalizedFinding
from src.schemas.segmentation import MaskResult


class StructuredReport(BaseModel):
    """Seven-field structured radiology report produced by Stage 3."""

    study_type: str = ""
    summary: str = ""
    main_findings: list[str] = Field(default_factory=list)
    detail_findings: list[str] = Field(default_factory=list)
    impression: str = ""
    recommendations: str = ""
    additional_informations: str = ""

    @field_validator("main_findings", "detail_findings", mode="before")
    @classmethod
    def coerce_to_list(cls, v: Any) -> list[str]:
        if isinstance(v, str):
            return [v] if v.strip() else []
        if isinstance(v, list):
            return [str(item) for item in v if str(item).strip()]
        return []


class ReportRequest(BaseModel):
    """All inputs needed to generate a radiology report (Stage 3)."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    case_id: str | None = None
    input_images: list[str]
    images: list[Any] = Field(default_factory=list, exclude=True)
    findings: list[LocalizedFinding] = Field(default_factory=list)
    masks: list[MaskResult] = Field(default_factory=list)


class ReportResult(BaseModel):
    case_id: str | None = None
    report_text: str                          # plain-text version (for PDF / fallback)
    structured: StructuredReport | None = None  # parsed structured fields
    status: Literal["success", "failed"]
    error: str | None = None


class PipelineOutput(BaseModel):
    """Final output of the three-stage CXR pipeline."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    detection: Any
    masks: list[MaskResult] = Field(default_factory=list)
    report: ReportResult | None = None
    processed_images: list[Any] = Field(default_factory=list, exclude=True)