app.py

from __future__ import annotations

import json
import traceback
from pathlib import Path
from typing import Any

import gradio as gr
from dotenv import load_dotenv
from PIL import Image

load_dotenv(Path(__file__).resolve().parent / ".env")

# ---------------------------------------------------------------------------
# Shared model store — heavy weights loaded once, reused across requests
# ---------------------------------------------------------------------------
_store: dict[str, Any] = {}

# ---------------------------------------------------------------------------
# Auto-initialize models at startup (handles HuggingFace Spaces deployment).
# startup.py downloads the MedSAM checkpoint and loads both models.
# It also updates os.environ["MEDSAM_CHECKPOINT"] so the config import below
# picks up the resolved path.
# ---------------------------------------------------------------------------
print("[app] Initializing AI models…")
_startup_status: str
try:
    from startup import initialize_all_models as _init_models
    _startup_status = _init_models(_store)
    print("[app] Model initialization complete.")
except Exception as _startup_err:
    _startup_status = (
        f"<span style='color:#C62828;font-weight:600'>"
        f"❌ Startup error: {_startup_err}</span>"
    )
    print(f"[app] Startup error: {_startup_err}")

# ---------------------------------------------------------------------------
# Config — imported AFTER startup so MEDSAM_CHECKPOINT env var is already set
# ---------------------------------------------------------------------------
from src.config.endpoints import (
    DEFAULT_MAX_FINDINGS,
    MEDGEMMA_MODEL_ID,
    MEDSAM_CHECKPOINT,
    MEDSAM_DEVICE,
)
from src.visualization import draw_segmentation_overlay, draw_view_localizations

# ---------------------------------------------------------------------------
# Visualization helpers
# ---------------------------------------------------------------------------

_STATUS_ICON: dict[str, str] = {
    "localized": "✅",
    "abstained": "⬜",
    "rejected_by_quality_gate": "⚠️",
    "parser_error": "❌",
}

_CERTAINTY_COLOR: dict[str, str] = {
    "positive": "🔴",
    "probable": "🟠",
    "questionable": "🟡",
}


def _findings_to_markdown(detection: Any) -> str:
    if not detection.findings:
        return "_No positive radiographic findings detected._"

    lines: list[str] = []
    for i, f in enumerate(detection.findings, 1):
        certainty_icon = _CERTAINTY_COLOR.get(f.certainty, "⚪")
        lines.append(f"### [{i}] {f.finding} {certainty_icon}")
        lines.append(f"- **Location:** {f.anatomical_location}")
        lines.append(f"- **Certainty:** {f.certainty}")
        lines.append("")
        lines.append("| View | Status | BBox |")
        lines.append("|------|--------|------|")
        for v in f.localizations:
            icon = _STATUS_ICON.get(v.status, "❓")
            if v.boxes:
                bbox_str = " · ".join(str(b.box_2d) for b in v.boxes)
            elif v.candidate_boxes:
                bbox_str = f"~~{v.candidate_boxes[0].box_2d}~~ (rejected)"
            else:
                bbox_str = "—"
            lines.append(f"| View {v.image_index} | {icon} `{v.status}` | {bbox_str} |")
        lines.append("")

    return "\n".join(lines)


def _structured_report_to_md(r: Any) -> str:
    """Render a StructuredReport as doctor-friendly markdown."""
    parts: list[str] = []

    if r.study_type:
        parts.append(f"## {r.study_type}\n")

    if r.summary:
        parts.append(f"**Summary**\n\n{r.summary}\n")

    if r.main_findings:
        items = "\n".join(f"- {f}" for f in r.main_findings)
        parts.append(f"**Main Findings**\n\n{items}\n")

    if r.detail_findings:
        items = "\n".join(f"- {f}" for f in r.detail_findings)
        parts.append(f"**Detailed Findings**\n\n{items}\n")

    if r.impression:
        parts.append(f"**Impression**\n\n{r.impression}\n")

    if r.recommendations:
        parts.append(f"**Recommendations**\n\n{r.recommendations}\n")

    ai = r.additional_informations or ""
    if ai.strip().lower() not in ("none", "n/a", "-", ""):
        parts.append(f"**Additional Information**\n\n{ai}\n")

    return "\n---\n\n".join(parts) if parts else "_No report generated._"


# ---------------------------------------------------------------------------
# Core functions called by Gradio
# ---------------------------------------------------------------------------


def run_analysis(
    img1: Image.Image | None,
    img2: Image.Image | None,
    medsam_ckpt: str,
    medsam_device: str,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[list, str, str, list, str]:
    """Return (annotated_gallery, findings_md, findings_json, seg_gallery, report_text)."""

    if "medgemma" not in _store:
        msg = "⚠️ AI system not initialized. Click **Initialize AI System** first."
        return [], msg, "{}", [], ""

    images = [img for img in [img1, img2] if img is not None]
    if not images:
        return [], "⚠️ Please upload at least one chest X-ray image.", "{}", [], ""

    try:
        from src.clients.medgemma_client import MedGemmaDetector, MedGemmaReporter
        from src.pipeline import CXRPipeline

        medgemma = _store["medgemma"]

        # Use segmentation model if available; skip silently if not loaded
        segmenter = _store.get("medsam")
        if segmenter is None and "medsam" not in _store:
            # Fallback: try to load (prefer startup-resolved path over UI textbox)
            try:
                progress(0.05, desc="Loading segmentation model…")
                from src.clients.medsam_client import MedSAMClient
                ckpt = _store.get("medsam_ckpt_path") or medsam_ckpt
                _store["medsam"] = MedSAMClient(checkpoint_path=ckpt, device=medsam_device)
                segmenter = _store["medsam"]
            except Exception:
                segmenter = None

        detector = MedGemmaDetector(client=medgemma, max_findings=DEFAULT_MAX_FINDINGS)
        reporter = MedGemmaReporter(client=medgemma)
        pipeline = CXRPipeline(detector=detector, segmenter=segmenter, reporter=reporter)

        input_images = [f"view_{i}.png" for i in range(len(images))]

        # ── Stage 1: Detection & Localization ────────────────────────────────
        progress(0.1, desc="Stage 1 — Finding discovery…")
        detection, processed_images = detector.detect(images, input_images)

        # ── Annotated images (bboxes) ────────────────────────────────────────
        progress(0.4, desc="Stage 1 — Rendering annotations…")
        annotated: list[Image.Image] = [
            draw_view_localizations(img, i, detection.findings)
            for i, img in enumerate(processed_images)
        ]

        # ── Findings text ────────────────────────────────────────────────────
        findings_md = _findings_to_markdown(detection)
        findings_json = json.dumps(detection.model_dump(), indent=2)

        # ── Stage 2: Segmentation ────────────────────────────────────────────
        masks: list = []
        seg_images: list[Image.Image] = []
        if segmenter is not None:
            progress(0.55, desc="Stage 2 — Segmentation…")
            masks = pipeline._run_segmentation(detection, list(processed_images))
            if masks:
                progress(0.75, desc="Stage 2 — Rendering segmentation masks…")
                seg_images = draw_segmentation_overlay(list(processed_images), masks)
        if not seg_images:
            seg_images = list(annotated)

        # ── Stage 3: Report generation ───────────────────────────────────────
        report_md = ""
        report_text = ""
        structured_report = None
        if reporter is not None:
            from src.schemas.report import ReportRequest
            progress(0.85, desc="Stage 3 — Generating radiology report…")
            report_request = ReportRequest(
                case_id=None,
                input_images=input_images,
                images=list(processed_images),
                findings=detection.findings,
                masks=masks,
            )
            report_result = reporter.generate_report(report_request)
            if report_result:
                if report_result.status == "success":
                    report_text = report_result.report_text
                    structured_report = report_result.structured
                    report_md = _structured_report_to_md(structured_report) if structured_report else report_text
                else:
                    report_md = f"⚠️ Report generation failed:\n\n{report_result.error}"

        # ── Persist last results for PDF export ──────────────────────────────
        _store["last_originals"]         = list(processed_images)
        _store["last_annotated"]         = annotated
        _store["last_seg_images"]        = seg_images if masks else None
        _store["last_detection"]         = detection
        _store["last_masks"]             = masks or None
        _store["last_report"]            = report_text or None
        _store["last_report_structured"] = structured_report

        progress(1.0, desc="Done.")
        return annotated, findings_md, findings_json, seg_images, report_md

    except Exception as exc:
        err = f"❌ Pipeline error:\n{exc}\n\n{traceback.format_exc()}"
        return [], err, "{}", [], ""


def export_pdf():
    """Generate a PDF and return a gr.update that sets value + visibility together."""
    if "last_detection" not in _store:
        gr.Warning("Please run an analysis before exporting.")
        return gr.update(value=None, visible=False)

    try:
        from src.export import export_to_tempfile
        path = export_to_tempfile(
            original_images=_store["last_originals"],
            annotated_images=_store["last_annotated"],
            seg_images=_store.get("last_seg_images"),
            detection=_store["last_detection"],
            masks=_store.get("last_masks"),
            report_text=_store.get("last_report"),
            structured_report=_store.get("last_report_structured"),
        )
        return gr.update(value=path, visible=True)
    except Exception as exc:
        gr.Warning(f"PDF export failed: {exc}")
        return gr.update(value=None, visible=False)


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

_KALBE_DARK = "#1A5C38"
_KALBE_MID  = "#2E7D32"
_KALBE_LIME = "#6DB33F"

_CSS = """
.report-box { font-family: 'Georgia', serif; line-height: 1.75; padding: 4px 8px; }
.report-box h2 { color: #1A5C38; margin-bottom: 4px; }
.report-box strong { color: #2E7D32; }
.report-box hr { border-color: #C8E6C9; margin: 12px 0; }
.init-status { padding: 4px 0 2px; }
footer { display: none !important; }

/* ── Kalbe header banner ── */
.kalbe-header {
    background: linear-gradient(135deg, #1A5C38 0%, #2E7D32 65%, #558B2F 100%);
    border-radius: 12px;
    padding: 22px 28px 18px;
    margin-bottom: 4px;
}
.kalbe-header h1 {
    color: #ffffff !important;
    font-size: 1.75rem !important;
    font-weight: 700 !important;
    margin: 0 0 2px !important;
    letter-spacing: 0.01em;
}
.kalbe-header p {
    color: #A5D6A7 !important;
    font-size: 1rem !important;
    font-weight: 500 !important;
    margin: 0 0 10px !important;
}
.kalbe-header .pipeline-pills {
    display: flex;
    gap: 8px;
    flex-wrap: wrap;
    margin-top: 8px;
}
.kalbe-header .pill {
    background: rgba(255,255,255,0.15);
    color: #E8F5E9 !important;
    border-radius: 20px;
    padding: 3px 12px;
    font-size: 0.8rem;
    font-weight: 500;
    border: 1px solid rgba(255,255,255,0.25);
}

/* ── PDF download widget ── */
.pdf-download {
    border: 1.5px dashed #2E7D32 !important;
    border-radius: 8px !important;
    background: #F1F8F2 !important;
    margin-top: 4px;
}
"""

with gr.Blocks(
    title="Thorax Report Generation — Kalbe Digital Lab",
    theme=gr.themes.Soft(
        primary_hue=gr.themes.colors.green,
        secondary_hue=gr.themes.colors.emerald,
        neutral_hue=gr.themes.colors.gray,
    ),
    css=_CSS,
) as demo:

    # ── Header ───────────────────────────────────────────────────────────────
    gr.HTML("""
    <div class="kalbe-header">
        <h1>Thorax Report Generation</h1>
        <p>Kalbe Digital Lab</p>
        <div class="pipeline-pills">
            <span class="pill">Detection &amp; Localization</span>
            <span class="pill">Segmentation</span>
            <span class="pill">Radiology Report</span>
        </div>
    </div>
    """)

    # ── Main layout ──────────────────────────────────────────────────────────
    with gr.Row(equal_height=False):

        # ── LEFT: Inputs & controls ──────────────────────────────────────────
        with gr.Column(scale=1, min_width=300):

            gr.Markdown("### Chest X-Ray Images")
            img1 = gr.Image(
                label="PA View (required)",
                type="pil",
                height=220,
            )
            img2 = gr.Image(
                label="Lateral View (optional)",
                type="pil",
                height=220,
            )

            gr.Markdown("---")

            with gr.Accordion("Advanced settings", open=False):
                model_id_box = gr.Textbox(
                    value=MEDGEMMA_MODEL_ID,
                    label="Detection model ID",
                )
                medsam_ckpt_box = gr.Textbox(
                    value=MEDSAM_CHECKPOINT,
                    label="Segmentation checkpoint path",
                )
                medsam_device_box = gr.Textbox(
                    value=MEDSAM_DEVICE,
                    label="Segmentation device",
                )

            gr.Markdown("---")
            run_btn    = gr.Button("Analyze X-Ray",      variant="primary",   size="lg")
            export_btn = gr.Button("Download Report (PDF)", variant="secondary", size="sm")
            pdf_output = gr.File(
                label="PDF Report",
                visible=False,
                elem_classes=["pdf-download"],
            )

        # ── RIGHT: Outputs ───────────────────────────────────────────────────
        with gr.Column(scale=2):
            with gr.Tabs():

                with gr.Tab("Annotated"):
                    annotated_gallery = gr.Gallery(
                        label="Detected findings with bounding boxes",
                        columns=2,
                        height=500,
                        object_fit="contain",
                        show_label=True,
                    )

                with gr.Tab("Findings"):
                    findings_md = gr.Markdown(
                        value="_Run the pipeline to see findings._"
                    )
                    with gr.Accordion("Raw JSON", open=False):
                        findings_json = gr.Code(
                            language="json",
                            label="DetectionResult",
                            lines=20,
                        )

                with gr.Tab("Segmentation"):
                    seg_gallery = gr.Gallery(
                        label="Polygon segmentation overlay",
                        columns=2,
                        height=500,
                        object_fit="contain",
                        show_label=True,
                    )
                    gr.Markdown(
                        "_Segmentation masks appear here when Stage 2 is enabled and "
                        "at least one finding is localized._"
                    )

                with gr.Tab("Report"):
                    report_box = gr.Markdown(
                        value="_Run the analysis to generate a structured radiology report._",
                        elem_classes=["report-box"],
                    )

    # ── Event wiring ─────────────────────────────────────────────────────────
    run_btn.click(
        fn=run_analysis,
        inputs=[
            img1,
            img2,
            medsam_ckpt_box,
            medsam_device_box,
        ],
        outputs=[
            annotated_gallery,
            findings_md,
            findings_json,
            seg_gallery,
            report_box,
        ],
    )

    export_btn.click(
        fn=export_pdf,
        outputs=[pdf_output],
    )


# ---------------------------------------------------------------------------
# Entry-point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )