src/clients/medgemma_client.py

from __future__ import annotations

import json
from collections.abc import Sequence

import torch
from PIL import Image
from pydantic import ValidationError
from transformers import pipeline

from src.config.endpoints import (
    DEFAULT_MAX_FINDINGS,
    DISCOVERY_MAX_NEW_TOKENS,
    LOCALIZATION_MAX_NEW_TOKENS,
    MEDGEMMA_MODEL_ID,
    REPORT_MAX_NEW_TOKENS,
)
from src.interfaces.detector import Detector
from src.interfaces.reporter import ReportGenerator
from src.localization_quality import evaluate_localization_boxes
from src.parsing import (
    extract_finding_discovery_payload,
    extract_json_payload,
    strip_medgemma_thinking_trace,
)
from src.preprocessing import pad_image_to_square
from src.prompts import (
    FINDING_DISCOVERY_PROMPT,
    build_localization_prompt,
    build_report_prompt,
)
from src.schemas.detection import (
    BBox,
    DetectionResult,
    Finding,
    FindingDiscovery,
    LocalizedFinding,
    ViewLocalization,
)
from src.schemas.report import ReportRequest, ReportResult, StructuredReport


def _select_torch_dtype() -> torch.dtype:
    if torch.cuda.is_available():
        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.float32


# ---------------------------------------------------------------------------
# Low-level model wrapper
# ---------------------------------------------------------------------------

class MedGemmaClient:
    """Wraps the MedGemma transformers pipeline for image-text-to-text inference."""

    def __init__(self, model_id: str = MEDGEMMA_MODEL_ID) -> None:
        self.model_id = model_id
        self._pipe = pipeline(
            "image-text-to-text",
            model=model_id,
            model_kwargs={
                "dtype": _select_torch_dtype(),
                "device_map": "auto",
            },
        )

    def generate(
        self,
        *,
        images: Sequence[Image.Image],
        prompt: str,
        max_new_tokens: int,
        system_instruction: str = "You are an expert radiologist.",
    ) -> str:
        content: list[dict] = [{"type": "image", "image": img} for img in images]
        content.append({"type": "text", "text": prompt})

        messages = [
            {"role": "system", "content": [{"type": "text", "text": system_instruction}]},
            {"role": "user", "content": content},
        ]
        output = self._pipe(text=messages, max_new_tokens=max_new_tokens, do_sample=False)
        response: str = output[0]["generated_text"][-1]["content"]
        return strip_medgemma_thinking_trace(response)


# ---------------------------------------------------------------------------
# Stage 1 – Detection & Localization
# ---------------------------------------------------------------------------

class MedGemmaDetector(Detector):
    """Discovers and localizes radiographic findings using MedGemma (Stage 1)."""

    def __init__(
        self,
        client: MedGemmaClient,
        max_findings: int = DEFAULT_MAX_FINDINGS,
    ) -> None:
        self._client = client
        self.max_findings = max_findings

    # -- internal helpers ----------------------------------------------------

    def discover_findings(self, images: list[Image.Image]) -> FindingDiscovery:
        response = self._client.generate(
            images=images,
            prompt=FINDING_DISCOVERY_PROMPT,
            max_new_tokens=DISCOVERY_MAX_NEW_TOKENS,
        )

        print("\n=== RAW FINDING DISCOVERY RESPONSE ===")
        print(response)
        print("======================================\n")

        payload = extract_finding_discovery_payload(response)

        print("\n=== NORMALIZED FINDING DISCOVERY PAYLOAD ===")
        print(json.dumps(payload, indent=2, ensure_ascii=False))
        print("============================================\n")

        return FindingDiscovery.model_validate(payload)

    def localize_finding(
        self,
        image: Image.Image,
        finding: Finding,
        image_index: int | None = None,
    ) -> list[BBox]:
        response = self._client.generate(
            images=[image],
            prompt=build_localization_prompt(finding),
            max_new_tokens=LOCALIZATION_MAX_NEW_TOKENS,
        )

        print("\n=== RAW LOCALIZATION RESPONSE ===")
        print(f"Finding     : {finding.finding}")
        print(f"Location    : {finding.anatomical_location}")
        if image_index is not None:
            print(f"Image index : {image_index}")
        print(response)
        print("=================================\n")

        payload = extract_json_payload(response, list)
        return [BBox.model_validate(item) for item in payload]

    # -- public interface ----------------------------------------------------

    def detect(
        self,
        images: Sequence[Image.Image],
        input_images: list[str],
        case_id: str | None = None,
    ) -> tuple[DetectionResult, list[Image.Image]]:
        if not images:
            raise ValueError("At least one image is required.")
        if len(images) != len(input_images):
            raise ValueError(
                f"images and input_images must have the same length; "
                f"got {len(images)} and {len(input_images)}."
            )

        processed_images = [pad_image_to_square(img) for img in images]
        discovery = self.discover_findings(processed_images)

        localized_findings: list[LocalizedFinding] = []

        for finding in discovery.findings[: self.max_findings]:
            view_localizations: list[ViewLocalization] = []

            for image_index, processed_image in enumerate(processed_images):
                try:
                    candidate_boxes = self.localize_finding(
                        processed_image, finding, image_index=image_index
                    )
                except (ValueError, ValidationError) as error:
                    print(
                        f"[warning] Could not parse localization for "
                        f"'{finding.finding}' on view {image_index}: {error}"
                    )
                    view_localizations.append(
                        ViewLocalization(
                            image_index=image_index,
                            image_path=input_images[image_index],
                            status="parser_error",
                            boxes=[],
                            candidate_boxes=[],
                            rejection_reasons=[f"localization parsing failed: {error}"],
                        )
                    )
                    continue

                if not candidate_boxes:
                    view_localizations.append(
                        ViewLocalization(
                            image_index=image_index,
                            image_path=input_images[image_index],
                            status="abstained",
                            boxes=[],
                            candidate_boxes=[],
                            rejection_reasons=[],
                        )
                    )
                    continue

                gate_result = evaluate_localization_boxes(
                    finding_name=finding.finding,
                    boxes=candidate_boxes,
                )

                if gate_result.accepted_boxes:
                    view_localizations.append(
                        ViewLocalization(
                            image_index=image_index,
                            image_path=input_images[image_index],
                            status="localized",
                            boxes=gate_result.accepted_boxes,
                            candidate_boxes=candidate_boxes,
                            rejection_reasons=gate_result.rejection_reasons,
                        )
                    )
                else:
                    view_localizations.append(
                        ViewLocalization(
                            image_index=image_index,
                            image_path=input_images[image_index],
                            status="rejected_by_quality_gate",
                            boxes=[],
                            candidate_boxes=candidate_boxes,
                            rejection_reasons=gate_result.rejection_reasons,
                        )
                    )

            localized_findings.append(
                LocalizedFinding(
                    finding=finding.finding,
                    anatomical_location=finding.anatomical_location,
                    certainty=finding.certainty,
                    localizations=view_localizations,
                )
            )

        result = DetectionResult(
            case_id=case_id,
            input_images=input_images,
            findings=localized_findings,
        )
        return result, processed_images


# ---------------------------------------------------------------------------
# Stage 3 – Report Generation
# ---------------------------------------------------------------------------

def _structured_to_text(r: StructuredReport) -> str:
    """Convert a StructuredReport to a plain-text representation for PDF / fallback."""
    parts: list[str] = []
    if r.study_type:
        parts.append(f"STUDY TYPE:\n{r.study_type}")
    if r.summary:
        parts.append(f"SUMMARY:\n{r.summary}")
    if r.main_findings:
        parts.append("MAIN FINDINGS:\n" + "\n".join(f"• {f}" for f in r.main_findings))
    if r.detail_findings:
        parts.append("DETAILED FINDINGS:\n" + "\n".join(f"• {f}" for f in r.detail_findings))
    if r.impression:
        parts.append(f"IMPRESSION:\n{r.impression}")
    if r.recommendations:
        parts.append(f"RECOMMENDATIONS:\n{r.recommendations}")
    ai = r.additional_informations
    if ai and ai.strip().lower() not in ("none", "n/a", "-", ""):
        parts.append(f"ADDITIONAL INFORMATION:\n{ai}")
    return "\n\n".join(parts)


class MedGemmaReporter(ReportGenerator):
    """Generates a structured radiology report using MedGemma (Stage 3)."""

    def __init__(self, client: MedGemmaClient) -> None:
        self._client = client

    def generate_report(self, request: ReportRequest) -> ReportResult:
        try:
            prompt = build_report_prompt(request.findings, masks=request.masks or None)

            print("\n=== GENERATING RADIOLOGY REPORT ===")
            print(f"Case ID   : {request.case_id}")
            print(f"Findings  : {len(request.findings)}")
            print(f"Masks     : {len(request.masks)}")
            print(f"Images    : {len(request.images)}")
            print("===================================\n")

            raw = self._client.generate(
                images=request.images,
                prompt=prompt,
                max_new_tokens=REPORT_MAX_NEW_TOKENS,
                system_instruction=(
                    "You are an expert radiologist. "
                    "Produce ONLY a single raw JSON object — "
                    "no markdown, no prose, no code fences."
                ),
            )

            print("\n=== RAW REPORT RESPONSE ===")
            print(raw)
            print("===========================\n")

            # Parse JSON → StructuredReport
            try:
                payload = extract_json_payload(raw, dict)
                structured = StructuredReport.model_validate(payload)
            except Exception as parse_exc:
                print(f"[warning] Could not parse structured report JSON: {parse_exc}")
                # Graceful fallback: put the raw text in impression
                structured = StructuredReport(
                    study_type="Chest X-Ray",
                    summary="Report generated (unstructured fallback).",
                    impression=raw,
                )

            report_text = _structured_to_text(structured)

            print("\n=== STRUCTURED REPORT ===")
            print(report_text)
            print("=========================\n")

            return ReportResult(
                case_id=request.case_id,
                report_text=report_text,
                structured=structured,
                status="success",
            )

        except Exception as exc:
            return ReportResult(
                case_id=request.case_id,
                report_text="",
                structured=None,
                status="failed",
                error=str(exc),
            )